Official

E - 圧縮番号列の復元 / Restoration of Compressed Number Sequence Editorial by admin

GPT 5.4 High

概要

「今までに何種類の箱を使ったか」だけを状態に持つ DP で数え上げます。
圧縮番号の性質により、次に置くボールが「既出の箱に入るか」「新しい箱に入るか」は、使った箱の種類数 \(k\) だけで判定・計算できます。

考察

この問題の本質は、実際の箱番号列 \(B\) を直接復元しようとすると場合分けが多すぎることです。
\(B_i\)\(1\) から \(M\) まで取り得るので、素朴に試すと \(M^N\) 通りになり到底間に合いません。

重要な観察 1:圧縮番号は「初登場順」

圧縮番号列では、ある時点までに異なる箱がちょうど \(k\) 個登場しているなら、割り当て済みの圧縮番号は必ず \(1,2,\dots,k\) です。

たとえば、ここまでに 3 種類の箱が登場しているなら、

  • 圧縮番号 1 の箱
  • 圧縮番号 2 の箱
  • 圧縮番号 3 の箱

がそれぞれちょうど 1 つずつ存在します。

つまり、「今までに何種類の箱が出たか」さえ分かれば、

  • 既にある圧縮番号 \(x\) を使うときは、その箱は ただ 1 つに決まる
  • 新しい圧縮番号を作るときは、まだ使っていない実際の箱番号から選ぶ

という形になります。


重要な観察 2:状態は「使った箱の種類数」だけで十分

\(f_i[k]\)

  • 先頭から \(i\) 個見たとき
  • 条件を満たし
  • 実際に使った箱番号の種類数がちょうど \(k\)

であるような箱番号列の個数とします。

この状態で十分な理由は、次の 1 手の選び方が \(k\) だけで決まるからです。

\(D_i = 0\)(圧縮番号が不明)のとき

次のボールは

  • 既に使った \(k\) 個の箱のどれかに入れる
    \(\Rightarrow\) \(k\) 通り
  • まだ使っていない箱に入れる
    \(\Rightarrow\) \(M-k\) 通り

のどちらでもよいです。

\(D_i = x \neq 0\)(圧縮番号が指定されている)のとき

次の圧縮番号は \(x\) でなければいけません。

  • すでに \(k \ge x\) 種類の箱を使っているなら、圧縮番号 \(x\) の箱は 1 つに決まる
    \(\Rightarrow\) 1 通り
  • まだ \(x-1\) 種類しか使っていないなら、今回初めて圧縮番号 \(x\) が登場する必要がある
    \(\Rightarrow\) 新しい箱を選ぶ必要があり、\(M-(x-1)\) 通り
  • \(x > k+1\) なら不可能
    (1 回で新しい圧縮番号は高々 1 つしか増えない)

ここがこの問題の肝です。
「実際にどの箱番号が圧縮番号 2 なのか」などを細かく持たなくても、個数だけなら \(k\) だけで遷移できる のです。


素朴な DP が難しい理由

たとえば「今まで使った実際の箱番号の集合」や「圧縮番号と実箱番号の対応」まで状態に持つと、状態数が爆発します。
しかし、上の観察から必要なのは「何種類使ったか」だけです。これで状態数は \(O(\min(N,M))\) に落ちます。

アルゴリズム

DP の定義

\(f_i[k]\) を「先頭 \(i\) 個まで見て、条件を満たし、使った箱の種類数がちょうど \(k\)」の個数とします。

初期値は

  • \(f_0[0] = 1\)
  • それ以外は \(0\)

です。


遷移 1:\(D_i = 0\) のとき

圧縮番号は不明なので、既存の箱を使っても新しい箱を使ってもよいです。

最終的に種類数が \(k\) になる方法は 2 つあります。

  1. もともと \(k\) 種類使っていて、今回もその中のどれかを使う
    \(k\) 通り
  2. もともと \(k-1\) 種類使っていて、今回新しい箱を使う
    → まだ未使用の箱は \(M-(k-1)=M-k+1\)

したがって

\[ f_i[k] = k \cdot f_{i-1}[k] + (M-k+1)\cdot f_{i-1}[k-1] \]

です。


遷移 2:\(D_i = x \neq 0\) のとき

今回の圧縮番号は \(x\) でなければいけません。

\(k < x\)

圧縮番号 \(x\) を出すには、少なくとも \(x\) 種類の箱が必要なので不可能です。

\(k = x\)

2 通りあります。

  1. すでに \(x\) 種類使っていて、そのうち圧縮番号 \(x\) の箱を使う
    → 1 通り
  2. もともと \(x-1\) 種類使っていて、今回新しい箱を導入してそれが圧縮番号 \(x\) になる
    \(M-x+1\) 通り

よって

\[ f_i[x] = f_{i-1}[x] + (M-x+1)\cdot f_{i-1}[x-1] \]

\(k > x\)

すでに圧縮番号 \(x\) の箱は存在しているので、それを使うしかありません。
したがって

\[ f_i[k] = f_{i-1}[k] \quad (k > x) \]

です。


不可能判定

指定された値 \(x=D_i\)

  • \(x > M\)
    (箱は \(M\) 種類しかない)
  • 現在可能な最大種類数を \(\text{lim}\) として、\(x > \text{lim}+1\)
    (1 回で新しい圧縮番号は 1 つしか増えない)

なら、その時点で答えは \(0\) です。


実装上の持ち方

コードでは 2 次元 DP 全体を持たず、現在の行だけを 1 次元配列 dp で持っています。

  • dp[k] : 現在の prefix で種類数が \(k\) の個数
  • lim : 現在到達可能な種類数の最大値

\(i\) ごとに新しい配列 new を作って更新します。
最後に sum(dp) を取れば、種類数が何で終わってもよいので答えになります。

計算量

  • 時間計算量: \(O(N \cdot \min(N,M))\)
  • 空間計算量: \(O(\min(N,M))\)

制約に \(N \times \min(N,M) \le 2.5 \times 10^7\) があるので、この DP で十分間に合います。

実装のポイント

  • 初期状態は「まだ何も見ていないので 0 種類」で、dp[0] = 1 にする。

  • D_i = x のとき、x > M または x > lim + 1 なら即座に 0 を出力する。

  • lim は「到達可能かどうか」を表すために別で持つ。
    DP の値は法 \(998244353\) で持っているので、たまたま 0 になっても「不可能」とは限らないためです。

  • D_i = 0lim = M のときは、これ以上新しい箱は増やせないことに注意する。

    ソースコード

import sys

MOD = 998244353

def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    N, M = data[0], data[1]
    D = data[2:2 + N]

    dp = [1]   # dp[k]: count for current prefix using exactly k distinct boxes
    lim = 0    # maximum reachable k (existence-wise, not modulo)

    mod = MOD
    m = M

    for d in D:
        if d == 0:
            prev = dp
            l = lim
            if l < m:
                new = [0] * (l + 2)
                for k in range(1, l + 1):
                    new[k] = (prev[k] * k + prev[k - 1] * (m - k + 1)) % mod
                new[l + 1] = prev[l] * (m - l) % mod
                dp = new
                lim = l + 1
            else:
                new = [0] * (l + 1)
                for k in range(1, l + 1):
                    new[k] = (prev[k] * k + prev[k - 1] * (m - k + 1)) % mod
                dp = new
        else:
            x = d
            if x > m or x > lim + 1:
                print(0)
                return

            prev = dp
            l = lim
            if x <= l:
                new = [0] * (l + 1)
                new[x:l + 1] = prev[x:l + 1]
                new[x] = (new[x] + prev[x - 1] * (m - x + 1)) % mod
                dp = new
            else:
                # x == l + 1
                new = [0] * (x + 1)
                new[x] = prev[l] * (m - l) % mod
                dp = new
                lim = x

    print(sum(dp) % mod)

if __name__ == "__main__":
    main()

この解説は gpt-5.4-high によって生成されました。

posted:
last update: