Official

N - 株価の補正 / Stock Price Correction Editorial 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 等張回帰は、以下の貪欲+ヒープアルゴリズムで解けます:

  1. 最大ヒープ(max-heap)を用意し、コスト total_cost = 0 とする。
  2. \(A_i\)\(= H_i - (i+1)\)、1-indexed のオフセット)を順に処理する。
  3. \(A_i\) について:
    • \(A_i\) をヒープに push する。
    • ヒープの最大値 \(\text{top}\) を確認する。
    • もし \(\text{top} > A_i\) ならば、非減少制約に違反しているので:
      • \(\text{top}\) を pop し、代わりに \(A_i\) をもう一度 push する(中央値を更新するイメージ)。
      • コストに \(\text{top} - A_i\) を加算する。
  4. 最終的な 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 によって生成されました。

posted:
last update: