Official

E - 石飛びの小道 / Stepping Stones Path Editorial by admin Editorial for printing

gemini-3.5-flash-thinking

概要

この問題は、飛び石 \(1\) から \(N\) まで移動する際に、最大 \(K\) 回まで「1つの石を飛び越える(立ち寄らない)」ことができるとき、訪れた石のスコアの合計を最大化する問題です。

「立ち寄らない石」の性質に着目して問題を言い換え、WQSバイナリサーチ(Alien DP)と呼ばれる高度な最適化手法を用いることで、制約を満たす高速な解法を導きます。


考察

1. 問題の言い換え(補集合を考える)

高橋君は「一歩進む」か「一つ飛ばしで跳ぶ」のいずれかを行えます。 「一つ飛ばしで跳ぶ」という行動は、「ある石を1つだけ飛ばして(立ち寄らずに)進む」と言い換えることができます。

ここで、飛ばす石の満たすべき条件を整理してみましょう。 - スタートの石 \(1\) とゴールの石 \(N\) には必ず立ち寄るため、飛ばせる石は \(2, 3, \ldots, N-1\) のみです。 - 「二つ飛ばし」以上の移動はできないため、連続する2つの石を同時に飛ばすことはできません。つまり、飛ばす石は互いに隣り合っていない(独立している)必要があります。 - 飛ばす石の個数(=「一つ飛ばしで跳ぶ」回数)は最大 \(K\) 回です。

すべての石のスコアの総和を \(S_{total}\) とし、飛ばした石のスコアの総和を \(T\) とすると、高橋君が訪れた石のスコアの合計は \(S_{total} - T\) となります。 \(S_{total}\) は定数なので、訪れた石のスコアの合計を最大化することは、「飛ばした石のスコアの総和 \(T\) を最小化すること」と同値です。

したがって、問題は以下のようにシンプルに再定義できます。

配列 \(B = [A_2, A_3, \ldots, A_{N-1}]\) から、互いに隣り合わない要素を最大 \(K\) 個選び、その総和を最小化する。

2. なぜ素直なDP(動的計画法)では間に合わないか?

この再定義された問題は、以下のDPで解くことができます。 - \(dp[i][j][0]\) : \(i\) 番目までの要素から \(j\) 個選び、 \(i\) 番目を選ばないときの最小値 - \(dp[i][j][1]\) : \(i\) 番目までの要素から \(j\) 個選び、 \(i\) 番目を選ぶときの最小値

しかし、このDPの状態数は \(O(N \times K)\) となります。 本問題の制約は \(N, K \leq 2 \times 10^5\) であるため、全体の計算量は \(O(NK) \approx 4 \times 10^{10}\) となり、実行時間制限(TLE)に引っかかってしまいます。

3. WQSバイナリサーチ(Alien DP)の導入

選ぶ個数 \(c\) に対する最小の総和を \(g(c)\) とします。選べる個数が増えるほど総和は小さくできる(または変わらない)ため、この関数 \(g(c)\) は \(c\) に関して下に凸な関数になります。

このように「選択個数に関する凸性」がある場合、WQSバイナリサーチ(Alien DP / ペナルティ法)が適用できます。

直感的なイメージ:ペナルティ(罰金)による個数制限のコントロール

「最大 \(K\) 個」という厳しい制限を直接扱う代わりに、「要素を1個選ぶ(石を1つ飛ばす)たびに、ペナルティ(罰金) \(C\) を支払う」という新しいルールを考えます。

  • 罰金 \(C\) が非常に大きい場合、誰も石を飛ばしたくないので、選ぶ個数は \(0\) 個になります。
  • 罰金 \(C\) が負の方向に非常に大きい(実質ボーナスがもらえる)場合、できるだけ多く飛ばしたいので、隣り合わない範囲で限界まで多くの石を選びます。

この罰金 \(C\) を二分探索でうまく調整することで、「罰金付きのルールで制限なしで最適化したときに、選んだ個数がちょうど \(K\) 個(付近)になるような絶妙な罰金額 \(C\)」を見つけることができます。

ペナルティ \(C\) を固定すれば、個数の制限がなくなるため、\(O(N)\) の単純なDPで解くことができます。


アルゴリズム

ステップ 1: 例外処理

\(K = 0\) または \(N \leq 2\) の場合は、石を飛ばすことができない(または飛ばす必要がない)ため、すべての石のスコアの総和 \(S_{total}\) がそのまま答えになります。

ステップ 2: ペナルティ \(C\) を固定したときの \(O(N)\) DP

配列 \(B\)(長さ \(M = N-2\))に対して、ペナルティ \(C\) を加味した最小化DPを行います。 同じ最小コストを達成できるなら、選んだ個数を最大化するようにタイブレークを行います。

  • \(dp0\) : 現在の要素を選ばないときの(最小コスト, \(-\)選んだ個数)
  • \(dp1\) : 現在の要素を選ぶときの(最小コスト, \(-\)選んだ個数) (Pythonのタプル比較 (cost, -cnt) を用いることで、コストの最小化と個数の最大化を同時に行えます)

遷移は以下の通りです。 - 次の要素を選ばない場合:前のステップの \(dp0\) と \(dp1\) のうち、より良い(コストが小さい、同値なら個数が多い)方を採用する。 - 次の要素を選ぶ場合:前のステップで選んでいない状態 \(dp0\) からのみ遷移可能。コストに現在の要素の値 \(b\) とペナルティ \(C\) を加算し、個数を \(1\) 増やす。

ステップ 3: 二分探索による最適な \(C\) の探索

ペナルティ \(C\) の探索範囲を \([-2 \times 10^9, 2 \times 10^9]\) とし、二分探索を行います。 - ペナルティ \(C\) でDPを解き、選ばれた個数が \(K\) 以上であれば、ペナルティが足りない(もっと罰金を重くして選ぶ個数を減らすべき)なので、探索範囲の「下限」を上げます。 - 選ばれた個数が \(K\) 未満であれば、罰金が重すぎるので、「上限」を下げます。

ステップ 4: 答えの復元

二分探索で得られた最適なペナルティを \(best\_C\) とします。 \(best\_C\) のもとでDPを解いて得られた最小コストを \(val\) とすると、これは「ペナルティ \(best\_C \times K\)」が含まれた値です。 したがって、ペナルティを差し引いた本来の最小コストは \(g(K) = val - best\_C \times K\) となります。

最終的な答えは、 \(S_{total} - g(K)\) です。


計算量

  • 時間計算量: \(O(N \log (\max |A_i|))\) 二分探索のステップ数は、探索範囲の広さ(約 \(4 \times 10^9\))から \(\log_2(4 \times 10^9) \approx 32\) 回です。各ステップで \(O(N)\) のDPを呼び出すため、全体の計算量は約 \(32 \times N\) となり、実行時間制限に余裕で間に合います。
  • 空間計算量: \(O(N)\) 入力の配列 \(A\) および \(B\) を保持するメモリのみが必要です。DPテーブルは直前の状態のみを保持すればよいため、空間計算量は \(O(N)\)(DP部分だけなら \(O(1)\))に抑えられます。

実装のポイント

  • タイブレークの重要性: WQSバイナリサーチでは、目的の個数 \(K\) が、凸関数の直線部分(傾きが一定で、同じペナルティで複数の個数が最適になり得る部分)に位置することがあります。そのため、同じコストなら「個数が最大になる方」を選ぶようにDPを実装し、二分探索の条件を cnt >= K とすることで、境界上の正しい値を確実に復元できます。

  • PyPy向けの高速化: 提示されたコードでは、DPの内部ループで新しいオブジェクト(タプルなど)の生成を避けるため、dp0_val, dp0_cnt などのプリミティブな変数に展開して更新を行っています。これにより、Python/PyPyにおける実行速度が劇的に向上します。

    ソースコード

import sys

def main():
    # Fast I/O
    input = sys.stdin.read
    data = input().split()
    if not data:
        return
    N = int(data[0])
    K = int(data[1])
    A = [int(x) for x in data[2:]]
    
    total_sum = sum(A)
    if K == 0 or N <= 2:
        print(total_sum)
        return
        
    B = A[1:N-1]
    M = len(B)
    
    # DP solver with penalty C
    # Optimized for PyPy by localizing variables and avoiding object creation
    def solve(C, B=B, M=M):
        dp0_val = 0
        dp0_cnt = 0
        dp1_val = B[0] + C
        dp1_cnt = -1
        
        for i in range(1, M):
            b = B[i]
            if dp0_val < dp1_val or (dp0_val == dp1_val and dp0_cnt < dp1_cnt):
                next_dp0_val = dp0_val
                next_dp0_cnt = dp0_cnt
            else:
                next_dp0_val = dp1_val
                next_dp0_cnt = dp1_cnt
            
            next_dp1_val = dp0_val + b + C
            next_dp1_cnt = dp0_cnt - 1
            
            dp0_val = next_dp0_val
            dp0_cnt = next_dp0_cnt
            dp1_val = next_dp1_val
            dp1_cnt = next_dp1_cnt
            
        if dp0_val < dp1_val or (dp0_val == dp1_val and dp0_cnt < dp1_cnt):
            return dp0_val, -dp0_cnt
        else:
            return dp1_val, -dp1_cnt

    # Check if the constraint K is active
    val, cnt = solve(0)
    if cnt <= K:
        print(total_sum - val)
        return
        
    # Binary search for the penalty C
    low = -2 * 10**9 - 100
    high = 2 * 10**9 + 100
    best_C = low
    
    while low <= high:
        mid = (low + high) // 2
        val, cnt = solve(mid)
        if cnt >= K:
            best_C = mid
            low = mid + 1
        else:
            high = mid - 1
            
    val, cnt = solve(best_C)
    g_K = val - best_C * K
    print(total_sum - g_K)

if __name__ == '__main__':
    main()

この解説は gemini-3.5-flash-thinking によって生成されました。

posted:
last update: