公式

C - 隣接ペナルティ付き選択 / Selection with Adjacent Penalty 解説 by admin

Claude 4.5 Opus

概要

\(N\) 個の仕事から1つ以上を選び、連続する番号の仕事を両方選ぶとペナルティ \(K\) がかかる条件下で、利益(報酬の合計 − ペナルティの合計)を最大化する問題です。

考察

重要な気づき

この問題では、仕事 \(i\) を引き受けるかどうかの決定が、直前の仕事 \(i-1\) を引き受けたかどうかにのみ依存します。なぜなら、ペナルティは「連続する番号のペア」に対してのみ発生するからです。

素朴なアプローチの問題点

全ての仕事の選び方を列挙すると \(2^N\) 通りあり、\(N \leq 2 \times 10^5\) では到底間に合いません(TLE)。

解決策

動的計画法(DP) を使います。各仕事について「引き受ける/引き受けない」の2状態を管理し、前の仕事の状態から現在の状態を計算することで、効率的に解けます。

アルゴリズム

状態の定義

仕事 \(i\) まで考慮したとき、以下の3つの状態を管理します:

  • prev_not_take_none: 仕事 \(i\) を選ばず、まだ1つも仕事を選んでいない場合の最大利益
  • prev_not_take_some: 仕事 \(i\) を選ばず、1つ以上の仕事を選んでいる場合の最大利益
  • prev_take: 仕事 \(i\) を選んでいる場合の最大利益(必ず1つ以上選んでいる)

遷移

仕事 \(i\) について:

仕事 \(i\) を引き受けない場合: - curr_not_take_none = prev_not_take_none(まだ何も選んでいない状態を維持) - curr_not_take_some = max(prev_not_take_some, prev_take)(以前に何か選んでいる状態を引き継ぐ)

仕事 \(i\) を引き受ける場合: - curr_take = max(prev_not_take_none + A[i], prev_not_take_some + A[i], prev_take + A[i] - K) - 前の仕事を選んでいない場合:ペナルティなしで \(A_i\) を得る - 前の仕事を選んでいる場合:ペナルティ \(K\) を払って \(A_i\) を得る

具体例

\(N=3, K=5, A=[10, 3, 8]\) の場合:

仕事 not_take_none not_take_some take
初期 0 \(-\infty\) 10
i=1 0 10 max(0+3, \(-\infty\)+3, 10+3-5) = 8
i=2 0 10 max(0+8, 10+8, 8+8-5) = 18

答え: \(\max(10, 18) = 18\)(仕事1と仕事3を選ぶ)

計算量

  • 時間計算量: \(O(N)\)(各仕事について定数時間の処理)
  • 空間計算量: \(O(1)\)(状態は直前の3つの値のみ保持)

実装のポイント

  1. 「1つ以上選ぶ」制約の処理: 単純なDPでは何も選ばない状態も含まれるため、「まだ何も選んでいない状態」と「1つ以上選んでいる状態」を分けて管理する必要があります。

  2. 初期値の設定: prev_not_take_some の初期値を \(-\infty\) にすることで、最初は「1つ以上選んでいるが現在の仕事は選ばない」状態が存在しないことを表現しています。

  3. 空間の最適化: 配列全体を保持する代わりに、直前の状態のみを変数で管理することで、空間計算量を \(O(1)\) に抑えています。

    ソースコード

def solve():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    
    # dp[i][j] = i番目の仕事まで考慮したときの最大利益
    # j = 0: 仕事iを引き受けない
    # j = 1: 仕事iを引き受ける
    
    # 初期状態: 何も選ばない状態から始める
    # ただし最終的に1つ以上選ぶ必要がある
    
    INF = float('inf')
    
    # dp[0]: 仕事iを引き受けない場合の最大利益
    # dp[1]: 仕事iを引き受ける場合の最大利益
    
    # 最初は何も選んでいない状態
    # prev_not_take: 前の仕事を引き受けなかった場合の最大利益
    # prev_take: 前の仕事を引き受けた場合の最大利益
    
    prev_not_take = 0  # 何も選んでいない
    prev_take = A[0]   # 仕事0を選ぶ
    
    for i in range(1, N):
        # 仕事iを引き受けない場合
        # 前の状態がどちらでも良い
        curr_not_take = max(prev_not_take, prev_take)
        
        # 仕事iを引き受ける場合
        # 前の仕事を引き受けなかった場合: 追加コストなし
        # 前の仕事を引き受けた場合: 追加コストK
        curr_take = max(prev_not_take + A[i], prev_take + A[i] - K)
        
        prev_not_take = curr_not_take
        prev_take = curr_take
    
    # 最終的に1つ以上選ぶ必要がある
    # prev_takeは少なくとも1つ選んでいる
    # prev_not_takeは0個の可能性もある
    
    # 1つ以上選んだ場合の最大値を求める
    # prev_takeは必ず1つ以上選んでいる
    # prev_not_takeから1つ以上選んでいるケースを考える必要がある
    
    # 再度DPをやり直す
    # dp[i][j][k]: i番目まで考慮、j=仕事iを選ぶか、k=1つ以上選んだか
    
    # prev_not_take_none: 前の仕事を選ばず、まだ何も選んでいない
    # prev_not_take_some: 前の仕事を選ばず、1つ以上選んでいる
    # prev_take: 前の仕事を選んでいる(必ず1つ以上選んでいる)
    
    prev_not_take_none = 0
    prev_not_take_some = -INF
    prev_take = A[0]
    
    for i in range(1, N):
        # 仕事iを引き受けない場合
        curr_not_take_none = prev_not_take_none
        curr_not_take_some = max(prev_not_take_some, prev_take)
        
        # 仕事iを引き受ける場合
        curr_take = max(prev_not_take_none + A[i], prev_not_take_some + A[i], prev_take + A[i] - K)
        
        prev_not_take_none = curr_not_take_none
        prev_not_take_some = curr_not_take_some
        prev_take = curr_take
    
    # 1つ以上選んでいる状態の最大値
    ans = max(prev_not_take_some, prev_take)
    print(ans)

solve()

この解説は claude4.5opus によって生成されました。

投稿日時:
最終更新: