N - 株価の補正 / Stock Price Correction Editorial by admin
Gemini 3.0 Flash (Thinking)概要
与えられた \(N\) 個の株価データ \(H_i\) を、整数制約の下で「厳密に単調増加(\(H'_1 < H'_2 < \cdots < H'_N\))」になるよう最小コストで補正する問題です。この問題は、適切な変数の置き換えを行うことで、L1等張回帰(Isotonic Regression)と呼ばれる問題に帰着でき、スロープトリック(Slope Trick)という手法を用いて効率的に解くことができます。
考察
1. 条件の言い換え
「厳密に単調増加」という条件 \(H'_1 < H'_2 < \cdots < H'_N\) は、各 \(H'_i\) が整数であることを踏まえると、以下のように書き換えることができます。 $\(H'_1 \leq H'_2 - 1 \leq H'_3 - 2 \leq \cdots \leq H'_N - (N-1)\)$
ここで、\(B_i = H'_i - (i-1)\) とおくと、条件は \(B_1 \leq B_2 \leq \cdots \leq B_N\) となり、「広義の単調増加(非減少)」という扱いやすい形になります。 また、元のコスト関数は以下のように変形できます。 $\(\sum_{i=1}^{N} |H_i - H'_i| = \sum_{i=1}^{N} |(H_i - (i-1)) - (H'_i - (i-1))| = \sum_{i=1}^{N} |A_i - B_i|\)\( ただし、\)A_i = H_i - (i-1)$ です。
結局、この問題は「与えられた数列 \(A\) を、非減少な整数列 \(B\) に最小コストで補正する」という問題に整理されました。
2. 動的計画法(DP)からのアプローチ
次のような DP を考えます。 \(dp[i][x] = (i \text{ 日目まで補正し、 } B_i = x \text{ としたときの最小コスト})\)
遷移式は次のようになります。 \(dp[i][x] = |x - A_i| + \min_{y \leq x} dp[i-1][y]\)
しかし、\(x\) の取りうる値の範囲(株価の範囲)が非常に広いため、このままでは計算が終わりません。ここで、関数 \(f_i(x) = dp[i][x]\) の形に注目すると、これは凸関数になっています。このような「凸関数の形状を維持しながら更新していく」テクニックがスロープトリックです。
アルゴリズム
スロープトリックによる最適化
非減少制約 \(B_1 \leq B_2 \leq \cdots \leq B_N\) の下での L1 距離最小化は、以下の手順で計算できます。
- 最大ヒープを用意する。
- 各 \(A_i\) について以下の操作を行う:
- ヒープに \(A_i\) を 2 回追加する。
- ヒープから最大値 \(top\) を取り出す。
- 現在の最小コストに \((top - A_i)\) を加算する。
この操作がなぜ正しいかを直感的に説明すると、ヒープの中身は「凸関数 \(f_i(x)\) の傾きが変化する点(折れ曲がり点)」を保持しています。 - \(|x - A_i|\) を加えることは、傾きを \(A_i\) を境に変化させることに対応します。 - \(\min_{y \leq x}\) を取る操作(累積最小値)は、凸関数の右側(傾きが正の部分)を切り捨てて平坦にすることに対応します。
最大ヒープのトップが常に「関数が最小値をとる右端の点」を表しており、新しい \(A_i\) がその点より左側にある場合、関数の最小値が押し上げられるため、その差分をコストに加算していきます。
計算量
- 時間計算量: \(O(N \log N)\)
- 各要素についてヒープへの push/pop を定数回行うため、\(N\) 回のループで \(O(N \log N)\) となります。
- 空間計算量: \(O(N)\)
- 要素数 \(N\) のヒープを保持するため、\(O(N)\) となります。
実装のポイント
最大ヒープ: Python の
heapqモジュールは最小ヒープ(最小値を取り出す)です。最大ヒープとして使うためには、値を負数にして格納する(-A_i)などの工夫が必要です。インデックスの調整: \(A_i = H_i - (i-1)\) を計算する際、 \(0\)\(-indexed\) の場合は \(A_i = H_i - i\) となります。
整数性: スロープトリックの過程で現れる折れ曲がり点はすべて入力 \(A_i\) に由来するため、補正後の値 \(B_i\) が整数であるという制約は自動的に満たされます。
ソースコード
import heapq
import sys
def solve():
"""
株価の補正問題を解く関数。
補正後の株価 H'_i が H'_1 < H'_2 < ... < H'_N を満たすようにし、
コストの合計 sum |H_i - H'_i| を最小化する。
この条件は B_i = H'_i - (i-1) と置くと、
B_1 <= B_2 <= ... <= B_N かつコスト sum |(H_i - (i-1)) - B_i| と同値。
これは L1 等張回帰 (Isotonic Regression) 問題であり、スロープトリックを用いて
O(N log N) で解くことができる。
"""
# 標準入力からすべてのデータを読み込み、空白で分割
input_data = sys.stdin.read().split()
if not input_data:
return
it = iter(input_data)
try:
# 1行目の N を取得
n = int(next(it))
except StopIteration:
return
cost = 0
# Pythonのheapqは最小ヒープなので、最大ヒープとして使うために値を反転させて格納する
max_heap = []
for i in range(n):
try:
# 各日の株価 H_i を取得
h_i = int(next(it))
except StopIteration:
break
# A_i = H_i - i (0-indexed)
# B_i = H'_i - i とすると、H'_i < H'_{i+1} は B_i <= B_{i+1} と同値
a_i = h_i - i
# スロープトリックによる L1 等張回帰の更新
# 1. 目的関数に |x - a_i| を加算する。
# 2. 累積最小値 (prefix minimum) を取る操作(非減少制約の維持)を行う。
# これらは、a_i をヒープに2回追加し、最大値を1回取り出すことで実現できる。
heapq.heappush(max_heap, -a_i)
heapq.heappush(max_heap, -a_i)
# 現在のヒープ内での最大値(スロープが正に転じる点)を取り出す
top = -heapq.heappop(max_heap)
# 最小コストの更新
# 新しい a_i がこれまでの最大値より小さい場合、その差分だけ最小値が増加する
cost += top - a_i
# 結果を出力
sys.stdout.write(str(cost) + '\n')
if __name__ == '__main__':
solve()
この解説は gemini-3-flash-thinking によって生成されました。
posted:
last update: