Official

N - 株価の補正 / Stock Price Correction Editorial by kyopro_friends


この問題は slope trick と呼ばれるテクニックを使って解くことができます。

「補正後の株価 \(H_1',H_2',\ldots,H_N'\)​ が \( H_1' < H_2', < \ldots < H_N'\) を満たす」(狭義単調増加)という条件は、\(H_i,H_i'\)\(H_i-i,H_i'-i\) に置き換えることで、他の条件を保ったまま「補正後の株価 \(H_1',H_2',\ldots,H_N'\)​ が \( H_1' \leq H_2', \leq \ldots \leq H_N'\) を満たす」(広義単調増加)という条件に置き換えることができます。以降、置き換え後の問題を考えます。

次のような DP を考えます。

\(\mathrm{dp}[n][x]\) を「\(H_1',\ldots,H_n'\) が広義単調増加かつ \(H_n' \leq x\) を満たすときの、\(\sum_{i=1}^{n}|H_i-H_i'|\) の最小値」とします。このとき、遷移は

\(\mathrm{dp}[n][x]=\min_{y\leq x}(\mathrm{dp}[n-1][y]+|H_n-y|)\)

となります。ここで、\(\mathrm{dp}[n]\)\(x\) の関数だと考えることにします。

\(n=1\) のとき

\(\mathrm{dp}[1][x]=\begin{cases} H_1-x & x\leq H_1\\ 0 & x\geq H_1 \end{cases}\)

なので \(\mathrm{dp}[1]\) は、\(x\) の関数として、傾きが整数の直線からなる区分線形連続凸関数です。任意の n について\(|H_n-x|\) もそうであり、そのような関数の和もそう、さらにその累積minもそうです。よって、帰納法により、任意の \(n\) について、 \(\mathrm{dp}[n]\)\(x\) の関数として、傾きが整数の直線からなる区分線形連続凸関数です。

そのような関数は、傾きの変化点と最小値のみを保持することで様々な操作を効率よく行うことができます。このテクニックを slope trick と呼びます。

この問題において、\(\mathrm{dp}[n-1]\) から \(\mathrm{dp}[n]\) を求めるには「定数 \(a\) により \(|x-a|\) の形で与えられる関数を足す」「累積minをとる」の2つの操作ができれば十分です。これは、傾きの変化点を優先度付きキュー \(S\) で持ちながら「\(S\)\(a\) を 2 つ追加し、最大の要素 \(r\) を削除する。関数の最小値に \(r-a\) を加える」とすることで達成できます。(詳細な説明は maspy さんによる解説記事に譲ります)

よって以上より \(O(N\log N)\) でこの問題を解くことができました。

実装例 (C++)

#include<bits/stdc++.h>
using namespace std;

int main(){
  int n;
  cin >> n;
  vector<int> h(n);
  for(int i=0; i<n; i++) cin >> h[i];
  for(int i=0; i<n; i++){
    h[i]-=i;
  }

  priority_queue<int> q;
  long long ans = 0;
  for(int i=0; i<n; i++){
    q.push(h[i]);
    q.push(h[i]);
    ans += q.top() - h[i];
    q.pop();
  }
  cout << ans << endl;
}

実装例 (Python)

import heapq

N = int(input())
H = list(map(int, input().split()))
H = [h - i for i, h in enumerate(H)]

q = []
ans = 0
for h in H:
  heapq.heappush(q, -h)
  v = -heapq.heappushpop(q, -h)
  ans += v - h
print(ans)

posted:
last update: