Official

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 距離最小化は、以下の手順で計算できます。

  1. 最大ヒープを用意する。
  2. 各 \(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: