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つ以上選ぶ」制約の処理: 単純なDPでは何も選ばない状態も含まれるため、「まだ何も選んでいない状態」と「1つ以上選んでいる状態」を分けて管理する必要があります。
初期値の設定:
prev_not_take_someの初期値を \(-\infty\) にすることで、最初は「1つ以上選んでいるが現在の仕事は選ばない」状態が存在しないことを表現しています。空間の最適化: 配列全体を保持する代わりに、直前の状態のみを変数で管理することで、空間計算量を \(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 によって生成されました。
投稿日時:
最終更新: