Official

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

Gemini 3.0 Flash (Thinking)

Overview

Given \(N\) stock price data points \(H_i\), the problem asks to correct them at minimum cost so that they become strictly monotonically increasing (\(H'_1 < H'_2 < \cdots < H'_N\)) under integer constraints. By performing an appropriate variable substitution, this problem can be reduced to what is known as L1 Isotonic Regression, which can be efficiently solved using a technique called Slope Trick.

Analysis

1. Reformulating the Condition

The strictly monotonically increasing condition \(H'_1 < H'_2 < \cdots < H'_N\), given that each \(H'_i\) is an integer, can be rewritten as follows: $\(H'_1 \leq H'_2 - 1 \leq H'_3 - 2 \leq \cdots \leq H'_N - (N-1)\)$

By defining \(B_i = H'_i - (i-1)\), the condition becomes \(B_1 \leq B_2 \leq \cdots \leq B_N\), which is the more tractable form of “non-strictly monotonically increasing (non-decreasing).” Additionally, the original cost function can be transformed as follows: $\(\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|\)\( where \)A_i = H_i - (i-1)$.

Ultimately, the problem is reformulated as “correct a given sequence \(A\) to a non-decreasing integer sequence \(B\) at minimum cost.”

2. Approach via Dynamic Programming (DP)

Consider the following DP: \(dp[i][x] = (\text{minimum cost when correcting up to day } i \text{ with } B_i = x)\)

The transition is as follows: \(dp[i][x] = |x - A_i| + \min_{y \leq x} dp[i-1][y]\)

However, since the range of possible values for \(x\) (the range of stock prices) is extremely large, this computation cannot finish as is. If we focus on the shape of the function \(f_i(x) = dp[i][x]\), it turns out to be a convex function. The technique of “updating while maintaining the shape of a convex function” is precisely Slope Trick.

Algorithm

Optimization via Slope Trick

Minimizing the L1 distance under the non-decreasing constraint \(B_1 \leq B_2 \leq \cdots \leq B_N\) can be computed with the following procedure:

  1. Prepare a max-heap.
  2. For each \(A_i\), perform the following operations:
    • Add \(A_i\) to the heap twice.
    • Extract the maximum value \(top\) from the heap.
    • Add \((top - A_i)\) to the current minimum cost.

To intuitively explain why this operation is correct: the contents of the heap maintain “the points where the slope of the convex function \(f_i(x)\) changes (breakpoints).” - Adding \(|x - A_i|\) corresponds to changing the slope with \(A_i\) as the boundary. - Taking \(\min_{y \leq x}\) (prefix minimum) corresponds to truncating the right side of the convex function (the part with positive slope) and making it flat.

The top of the max-heap always represents “the rightmost point where the function attains its minimum value.” When a new \(A_i\) is to the left of that point, the minimum value of the function is pushed up, so the difference is added to the cost.

Complexity

  • Time Complexity: \(O(N \log N)\)
    • Since a constant number of push/pop operations on the heap are performed for each element, the \(N\) iterations result in \(O(N \log N)\).
  • Space Complexity: \(O(N)\)
    • Since we maintain a heap of \(N\) elements, the space is \(O(N)\).

Implementation Notes

  • Max-heap: Python’s heapq module is a min-heap (extracts the minimum value). To use it as a max-heap, a workaround such as storing negated values (-A_i) is needed.

  • Index adjustment: When computing \(A_i = H_i - (i-1)\), if using \(0\)-indexed, it becomes \(A_i = H_i - i\).

  • Integrality: Since all breakpoints that appear during the Slope Trick process originate from the input \(A_i\), the constraint that the corrected values \(B_i\) are integers is automatically satisfied.

    Source Code

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()

This editorial was generated by gemini-3-flash-thinking.

posted:
last update: