Official

F - ピアノの練習 Editorial by harurun4635


「鍵盤の選び方」を直接数えようとすると、各選び方について、実際にすべての鍵盤を同時に押せるかを判定しなければなりません。

そこで、「鍵盤の選び方」に対して「腕・指の押し方」を一意に定め、両者を \(1\)\(1\) に対応させます。これにより、鍵盤を選んだ後に可能かを判定するのではなく、条件を満たす押し方だけを DP で構築しながら数えられます。

いわゆる、数えたい対象に一意な標準形を与え、その標準形を代わりに数えるというテクニックです。


今回は、選んだ鍵盤を左から順に見て、腕・指を次のように割り当てます。

  • 最初の鍵盤は、\(1\) 本目の腕の \(1\) 本目の指で押す
  • 現在の腕で押せる(指が余っていて・ はじめの指と \(D\) 以下)なら、その腕の次の指で押す
  • 現在の腕で押せないなら、次の腕の \(1\) 本目の指で押す

この貪欲法で使う腕の本数は、選んだ鍵盤を押すために必要な腕の本数の最小値になります。(本来は証明が必要でしょうが、直感的に明らかなわりに煩雑になるので略させてください)


鍵盤の選び方を固定すると、上の貪欲法によって、どの腕・指で押すかが一意に決まります。逆に、この貪欲法による押し方を固定すると、押している鍵盤の集合も一意に決まります。

また、貪欲法が使う腕の本数は最小なので、選んだ鍵盤を同時に押せることと、貪欲法で使う腕が \(A\) 本以下であることは同値です。

よって、この押し方を DP で数えればよいです。


DP

鍵盤を \(1,2,\ldots,N\) の順に処理します。DP の状態を (j, a, s, b) とします。

  • j: 選んだ鍵盤の個数
  • a: 使った腕の本数
  • s: 現在の腕が最初に押した鍵盤
  • b: 現在の腕で押している鍵盤の個数

次に見る鍵盤を \(x\) とすると、次の遷移を行います。

  • 鍵盤 \(x\) を選ばない
  • 鍵盤 \(x\) を選ぶ
    • まだ鍵盤を選んでいないなら、新しい腕で押す
    • \(b<B\) かつ \(x-s\le D\) なら、現在の腕で押す
    • そうでなければ、新しい腕で押す

状態数は 状態数は \(O(N \times NKAB)\) 遷移は \(O(1)\) より計算量は \(O(N^2KAB)\) です。

陽に空間 \(O(N \times NKAB)\) とる方針では MLE に注意してください。例えば nextdp の形にすることで空間 \(O(NKAB)\) となります。


実装例

from collections import defaultdict

mod = 998244353

n, k, aa, bb, d = map(int, input().split())

dp = {(0, 0, 0, 0): 1}

for x in range(1, n + 1):
    ndp = defaultdict(int)

    for (j, a, s, b), v in dp.items():
        # x を選ばない
        ndp[j, a, s, b] += v

        if j == k:
            continue

        # x を選ぶ
        if b == 0:
            if a < aa:
                ndp[j + 1, a + 1, x, 1] += v
        elif b < bb and x - s <= d:
            ndp[j + 1, a, s, b + 1] += v
        elif a < aa:
            ndp[j + 1, a + 1, x, 1] += v

    dp = {key: v % mod for key, v in ndp.items()}

ans = 0
for (j, a, s, b), v in dp.items():
    if j == k:
        ans += v

print(ans % mod)

posted:
last update: