D - アルバイトの計画 / Part-Time Job Planning Editorial by admin Editorial for printing
Gemini 3.8 Flash (High)概要
\(N\) 日間のうち、一度アルバイトをしたらその後 \(D\) 日間はアルバイトができないという制約のもとで、最大 \(K\) 日アルバイトをしたときに得られる給料の合計の最大値を求める問題です。
考察
\(D = 0\) の場合
休養期間が不要(\(D = 0\))である場合、日付の間隔に関する制約はありません。したがって、給料 \(A_i\) が大きい順に上位 \(\min(N, K)\) 個の日を選んで働くのが最適です。これは配列のソートにより \(O(N \log N)\) で解くことができます。
\(D \ge 1\) の場合
ある日にアルバイトをすると直後の \(D\) 日間は働けなくなるため、「どの日を選んだか」によってその後に選べる日の選択肢が変わります。このように選んだ履歴が影響し、かつ未来の最適化に重なる部分問題が存在するため、動的計画法 (DP) を用いるのが自然です。
素朴に「どの日に働いたか」を全探索しようとすると、組み合わせの数が指数関数的に増加し、実行時間制限 (TLE) になってしまいます。
DP の状態と遷移の設計
日数を表す添字を \(i\)、アルバイトをした日数を表す添字を \(j\) とします。 素直に状態を考えると、 - \(DP[j][i] = \) 「\(i\) 日目までに合計 \(j\) 日アルバイトをしたときの給料の最大値」 と定義できます。
\(i\) 日目の行動として以下の 2 通りが考えられます: 1. \(i\) 日目にアルバイトをしない場合: \(i-1\) 日目までの最大値がそのまま引き継がれます(\(DP[j][i-1]\))。 2. \(i\) 日目にアルバイトをする場合: 直前にアルバイトができるのは \(i - (D + 1)\) 日目以前です。したがって、\(\max(0, i - (D + 1))\) 日目までに \(j-1\) 日アルバイトをしたときの最大値に \(A_i\) を加算したものになります(\(DP[j-1][\max(0, i - D - 1)] + A_i\))。
この 2 つの最大値をとることで、各マスを \(O(1)\) で遷移できます。
\(K\) の上限の削減
\(N\) 日間で \(D\) 日の間隔を空けながら選べる日数の最大値は、 \(1 + \left\lfloor \frac{N - 1}{D + 1} \right\rfloor\) 日 です。入力で与えられる \(K\) がこれより大きい場合、物理的に \(K\) 日働くことは不可能です。あらかじめ \(K\) をこの最大値で切り詰めておくことで、無駄な計算を削減できます(\(D \ge 1\) のとき、切り詰めた \(K\) は最大でも約 \(N / 2\) です)。
アルゴリズム
- コーナーケースの処理: \(D = 0\) の場合は、配列 \(A\) を降順にソートし、先頭から \(K\) 個の要素の和を出力して終了します。
- \(K\) の上限を切り詰める: \(K \leftarrow \min\left(K, 1 + \left\lfloor \frac{N - 1}{D + 1} \right\rfloor\right)\) と更新します。
- DPテーブルの更新:
- \(DP[j][i]\) の更新には \(j-1\) 行目の結果しか必要としないため、2行分(直前の行
prev_fと現在の行curr_f)の配列を用意して使い回します(ローリング配列)。 - \(j = 1\) から \(K\) まで順にループを回します。
- \(j\) 回目のアルバイトを行うためには最低でも \((j - 1)(D + 1) + 1\) 日目以降である必要があるため、内側のループ \(i\) の開始地点をこの値からにします。
- 各 \(i\) について遷移式に従って
curr_f[i]を計算します: $\(p = \max(0, i - (D + 1))\)\( \)\(\text{curr\_f}[i] = \max(\text{curr\_f}[i - 1], \text{prev\_f}[p] + A_i)\)$ - 各 \(j\) について最終日までの最大値 \(\text{curr\_f}[N]\) で答えを更新していきます。
- \(DP[j][i]\) の更新には \(j-1\) 行目の結果しか必要としないため、2行分(直前の行
計算量
- 時間計算量: \(O(N \cdot K)\)
- \(D = 0\) のときはソートの \(O(N \log N)\) です。
- \(D \ge 1\) のとき、切り詰め後の \(K\) は最大でも約 \(N/2\) となります。外側ループが \(K\) 回、内側ループが最大 \(N\) 回であるため、全体の計算量は \(O(N \cdot K) \le O(N^2 / 2)\) です。\(N = 5000\) のときループ回数は高々 \(1.25 \times 10^7\) 回程度であり、実行時間制限に十分間に合います。
- 空間計算量: \(O(N)\)
- 長さ \(N + 1\) の一次元配列を 2 つのみ保持するため、空間計算量は \(O(N)\) と非常に省メモリです。
実装のポイント
累積 \(\max\) を状態に組み込む: \(DP[j][i]\) を「\(i\) 日目にちょうど \(j\) 回目のアルバイトをしたときの最大値」ではなく、「\(i\) 日目までに \(j\) 回アルバイトをしたときの最大値」と定義することで、直前の最適値を \(O(1)\) で参照できるようにしています。
不可能な状態の初期化: アルバイトを行う条件を満たせない状態からの遷移を防ぐため、DPテーブルは十分小さな値(\(-10^{18}\) などの \(-\infty\))で初期化しておきます。
ソースコード
import sys
def main():
input_data = sys.stdin.read().split()
if not input_data:
return
N = int(input_data[0])
K = int(input_data[1])
D = int(input_data[2])
A = [int(x) for x in input_data[3 : 3 + N]]
if D == 0:
A.sort(reverse=True)
print(sum(A[:K]))
return
max_k = 1 + (N - 1) // (D + 1)
K = min(K, max_k)
INF = 10**18
prev_f = [0] * (N + 1)
ans = 0
step = D + 1
for j in range(1, K + 1):
curr_f = [-INF] * (N + 1)
start_i = (j - 1) * step + 1
for i in range(start_i, N + 1):
p = i - step if i > step else 0
val = prev_f[p] + A[i - 1]
c = curr_f[i - 1]
curr_f[i] = val if val > c else c
if curr_f[N] > ans:
ans = curr_f[N]
prev_f = curr_f
print(ans)
if __name__ == "__main__":
main()
この解説は gemini-3.8-flash-high によって生成されました。
posted:
last update: