概要

アルバイトをする日の間隔を \(D+1\) 日以上空けながら、最大 \(K\) 日を選び、給料の合計を最大化します。

「何日目までを考えたか」と「何日まで働けるか」を状態とする動的計画法（DP）を使い、特別な場合はより簡単な方法で処理します。

考察

給料の高い日から選ぶだけではうまくいかない

\(D>0\) の場合、ある日を選ぶと、その近くの日を選べなくなります。

例えば \(N=3,\ K=2,\ D=1\)、給料が \([6,10,6]\) のとき、最も給料が高い \(2\) 日目を選ぶと、合計は \(10\) 円です。一方、\(1\) 日目と \(3\) 日目を選べば、合計 \(12\) 円になります。

したがって、単純に給料の高い日を優先する貪欲法では正解できません。また、働く日の組合せをすべて調べる方法も、最大 \(2^N\) 通りあるため間に合いません。

最後の日に働くかどうかで分ける

\(F_k(i)\) を、次のように定義します。

• \(1\) 日目から \(i\) 日目までで、アルバイトをする日数を \(k\) 日以下 にしたときの給料の最大値

\(i\) 日目について、次の二択を考えます。

• 働かない場合：最大値は \(F_k(i-1)\)
• 働く場合：直前に働けるのは \(i-D-1\) 日目以前なので、最大値は \(F_{k-1}(\max(0,i-D-1))+A_i\)

よって、遷移は次のようになります。

\(F_k(i)=\max\left(F_k(i-1),\ F_{k-1}(\max(0,i-D-1))+A_i\right)\)

初期値は \(F_0(i)=0,\ F_k(0)=0\) です。

この方法なら、各状態を定数時間で計算でき、時間計算量は \(O(NK)\) になります。

働ける日数には、休養条件による上限がある

\(step=D+1\) とします。\(m\) 日働くためには、最も詰めて働いても、最後に働く日は \(1+(m-1)step\) 日目です。

したがって、働ける日数の最大値は、

\(M=\left\lfloor\frac{N+D}{D+1}\right\rfloor\)

です。

\(K\ge M\) なら、休養条件を守るだけで日数制限も満たします。この場合は、DP の状態から日数 \(k\) を取り除けます。

アルゴリズム

コードでは、次の \(3\) 通りに分けて処理しています。

1. \(D=0\) の場合

休養期間がないため、どの日も自由に選べます。

すべての給料は正なので、給料を降順に並べ、上位 \(K\) 個の合計を出力します。

2. \(K\ge M\) の場合

日数制限を考えなくてよいため、一次元 DP を使います。

\(dp[i]\) を「\(1\) 日目から \(i\) 日目までで得られる給料の最大値」とすると、

\(dp[i]=\max\left(dp[i-1],\ dp[\max(0,i-step)]+A_i\right)\)

です。

\(i=1,2,\ldots,N\) の順に計算し、\(dp[N]\) を出力します。

コードの best は、更新前には \(dp[i-1]\)、更新後には \(dp[i]\) を表しています。

3. それ以外の場合

考察で示した \(F_k(i)\) を計算します。ただし、すべての \(k\) の結果を保存する必要はありません。

\(F_k(i)\) の計算に必要なのは、

• 前の段階の値 \(F_{k-1}(i-step)\)
• 現在の段階で直前までに求めた値 \(F_k(i-1)\)

だけなので、配列を使い回してメモリを節約できます。

最初に、\(1\) 日以下働く場合を計算する

\(1\) 日しか働かない場合、休養条件は気にしなくてよいため、

\(F_1(i)=\max(A_1,A_2,\ldots,A_i)\)

です。これを dp に格納します。

\(k=2,3,\ldots,K\) の順に更新する

更新前の dp[i] は \(F_{k-1}(i)\) を表します。新しい結果は current に格納します。

ここで、\(k\) 日働ける最初の日は、

\(start=(k-1)step+1\)

です。\(i<start\) では、そもそも \(k\) 日働くことができないため、\(F_k(i)=F_{k-1}(i)\) です。

そこで、次のように更新します。

1. current = dp[:] として、前の段階をコピーする。
2. best = dp[start - 1] とする。
3. \(i=start,\ldots,N\) の順に処理する。
   • \(i\) 日目に働く場合の値 dp[i - step] + A[i] を求める。
   • その値と best の大きい方を、新しい best とする。
   • current[i] = best とする。
4. dp = current とする。

best が \(F_k(i-1)\) を保持するため、この更新は考察の遷移式そのものです。

最後に dp[N]、すなわち \(F_K(N)\) を出力します。

計算量

• 時間計算量:
  • \(D=0\) の場合：\(O(N\log N)\)
  • \(D>0\) かつ \(K\ge M\) の場合：\(O(N)\)
  • それ以外の場合：\(O(NK)\)
• 空間計算量: \(O(N)\)

実装のポイント

• 必要な間隔は \(D\) ではなく \(D+1\)

\(i\) 日目に働いた後、次に働けるのは \(i+D+1\) 日目です。step = D + 1 とすると、添字の計算が分かりやすくなります。

• DP は「ちょうど \(k\) 日」ではなく「\(k\) 日以下」

休養条件によっては、給料の高い日を少数だけ選ぶ方が有利な場合もあります。日数を必ず使い切る必要はありません。

• 前の段階の配列を直接書き換えない

dp[i - step] には、あくまで \(k-1\) 日以下の結果が必要です。current を別に用意して、更新済みの値が混ざらないようにします。

• 答えには大きな整数が必要

給料の合計は最大 \(5000\times10^9=5\times10^{12}\) になり得ます。Python ではそのまま扱えますが、他の言語では 64 ビット整数を使います。

• \(D=0\) の分岐では、先頭のダミー要素もソートされる

コードの A には添字調整用の \(0\) が含まれていますが、すべての \(A_i\) は正で \(K\le N\) なので、降順の上位 \(K\) 個にこの \(0\) が入ることはありません。

ソースコード

import sys


def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    N, K, D = data[:3]
    A = [0] + data[3:]

    if D == 0:
        A.sort(reverse=True)
        print(sum(A[:K]))
        return

    step = D + 1

    if K >= (N + D) // step:
        dp = [0] * (N + 1)
        best = 0
        for i in range(1, N + 1):
            value = A[i]
            if i > step:
                value += dp[i - step]
            if value > best:
                best = value
            dp[i] = best
        print(best)
        return

    dp = [0] * (N + 1)
    best = 0
    for i in range(1, N + 1):
        if A[i] > best:
            best = A[i]
        dp[i] = best

    for k in range(2, K + 1):
        start = (k - 1) * step + 1
        current = dp[:]
        best = dp[start - 1]
        for i in range(start, N + 1):
            value = dp[i - step] + A[i]
            if value > best:
                best = value
            current[i] = best
        dp = current

    print(dp[N])


if __name__ == "__main__":
    main()

---

この解説は gpt-6-astra-medium によって生成されました。