N - 株価の補正 / Stock Price Correction 解説 by admin
Claude 4.6 Opus (Thinking)概要
\(N\) 日間の株価データを狭義単調増加に補正するときの、変更コスト(絶対値の総和)の最小値を求める問題です。これは「L1ノルムでの等張回帰(isotonic regression)」を Slope Trick で高速に解くことができます。
考察
狭義単調増加から非減少への変換
補正後の株価 \(H'_i\) が狭義単調増加(\(H'_1 < H'_2 < \cdots < H'_N\))であることを求められています。ここで、\(A_i = H_i - i\) および \(B_i = H'_i - i\) と変換すると、
\[H'_1 < H'_2 < \cdots < H'_N \iff B_1 \leq B_2 \leq \cdots \leq B_N\]
が成り立ちます(整数制約があるため、差が \(1\) 以上 \(\Leftrightarrow\) \(i\) を引けば差が \(0\) 以上)。またコストは
\[\sum_{i=1}^{N} |H_i - H'_i| = \sum_{i=1}^{N} |A_i - B_i|\]
と変わりません。これにより、問題は 「\(B\) を非減少列にしたとき \(\sum |A_i - B_i|\) を最小化する」 という L1 ノルム等張回帰問題に帰着されます。
素朴なアプローチの問題点
DP で「\(B_i\) の取りうる値」を全探索すると、値の範囲が最大 \(10^9\) に達するため到底間に合いません。座標圧縮して \(O(N^2)\) にしても \(N \leq 2 \times 10^5\) では TLE です。
Slope Trick による解法
L1 ノルムの最小化では、コスト関数が区分線形凸関数になることを利用します。Slope Trick と呼ばれるテクニックでは、この区分線形凸関数の「傾きが変わる点(折れ点)」を優先度付きキュー(ヒープ)で管理することで、効率的に更新できます。
アルゴリズム
非減少列への L1 等張回帰は、以下の貪欲+ヒープアルゴリズムで解けます:
- 最大ヒープ(max-heap)を用意し、コスト
total_cost = 0とする。 - \(A_i\)(\(= H_i - (i+1)\)、1-indexed のオフセット)を順に処理する。
- 各 \(A_i\) について:
- \(A_i\) をヒープに push する。
- ヒープの最大値 \(\text{top}\) を確認する。
- もし \(\text{top} > A_i\) ならば、非減少制約に違反しているので:
- \(\text{top}\) を pop し、代わりに \(A_i\) をもう一度 push する(中央値を更新するイメージ)。
- コストに \(\text{top} - A_i\) を加算する。
- 最終的な
total_costが答え。
直感的な説明: ヒープには「これまでの最適解における中央値的な代表点」が格納されています。新しい値が現在の代表点より小さい場合、非減少制約を守るために調整が必要で、その調整コストが \(\text{top} - A_i\) です。ヒープの更新(pop して push)は、過去の区間の最適な代表値を下方修正する操作に対応します。
具体例
\(N = 3\), \(H = [5, 3, 4]\) の場合:
- \(A = [5 - 1, 3 - 2, 4 - 3] = [4, 1, 1]\)
- \(i=0\): push 4 → heap = [4], top = 4, \(4 \leq 4\) なので何もしない。cost = 0
- \(i=1\): push 1 → heap = [4, 1], top = 4, \(4 > 1\) なので pop 4, push 1 → heap = [1, 1], cost += 3 → cost = 3
- \(i=2\): push 1 → heap = [1, 1, 1], top = 1, \(1 \leq 1\) なので何もしない。cost = 3
答えは 3。実際、\(H' = [3, 4, 5]\)(\(B = [2, 2, 2]\))とすれば \(|5-3|+|3-4|+|4-5| = 2+1+1 = 4\)… ではなく、\(H' = [2, 3, 4]\)(\(B = [1, 1, 1]\))とすれば \(|5-2|+|3-3|+|4-4| = 3\) でコスト 3 が達成されます。
計算量
- 時間計算量: \(O(N \log N)\)(各要素についてヒープ操作が \(O(\log N)\))
- 空間計算量: \(O(N)\)(ヒープのサイズが最大 \(N\))
実装のポイント
Python の
heapqは最小ヒープなので、値の符号を反転して格納することで最大ヒープとして使います。狭義単調増加 → 非減少への変換で使うオフセットは \(i+1\)(1-indexed)です。\(i\) を 0-indexed で回す場合は \(A_i = H_i - (i+1)\) とします。
\(H_i\) や \(A_i\) の値が大きくなりうるので、コストの合計が 64bit 整数の範囲に収まるか注意が必要ですが、Python では整数のオーバーフローは起こりません。
ソースコード
import heapq
def solve():
import sys
input_data = sys.stdin.buffer.read().split()
N = int(input_data[0])
H = [int(input_data[i+1]) for i in range(N)]
# Transform to non-decreasing problem by subtracting index
# Strictly increasing H'[i] means H'[i] - i is non-decreasing
# Let A[i] = H[i] - i, then we want to find non-decreasing B[i] minimizing sum |A[i] - B[i]|
A = [H[i] - (i + 1) for i in range(N)] # using 1-indexed offset so H'_1 < H'_2 < ... < H'_N becomes B non-decreasing
# Slope trick for isotonic regression (L1 norm)
# We maintain the piecewise linear convex function f represented by a max-heap of slopes
# For non-decreasing sequence minimizing L1 cost:
# We use a max-heap. For each new element a:
# - push a into heap
# - if top of heap > a, we pop the top, push a again (merge operation), and add (top - a) to cost
max_heap = [] # stored as negatives for max-heap behavior
total_cost = 0
for a in A:
heapq.heappush(max_heap, -a)
top = -max_heap[0]
if top > a:
heapq.heappop(max_heap)
heapq.heappush(max_heap, -a)
total_cost += top - a
print(total_cost)
solve()
この解説は claude4.6opus-thinking によって生成されました。
投稿日時:
最終更新: