Official

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

Claude 4.6 Opus (Thinking)

概要

箱番号列を「圧縮番号列」に変換する手続きにおいて、一部が欠損した圧縮番号列の条件を満たす元の箱番号列の個数を数え上げる問題です。DPにより、現在までに登場した異なる箱番号の種類数を状態として管理して解きます。

考察

圧縮番号列の性質

圧縮番号列 \(C_1, C_2, \dots, C_N\) には重要な性質があります。圧縮番号は「初めて登場した箱番号」に対して \(1, 2, 3, \dots\) と順番に割り当てられるため:

  • \(C_i\)新しい値(初めて登場する圧縮番号)であるとき、\(C_i\) はちょうど「それまでに登場した異なる箱番号の種類数 \(+1\)」でなければなりません。
  • \(C_i\)既存の値であるとき、\(C_i\) は「それまでに既に登場した圧縮番号」のいずれかです。

例えば、\(i=1\)\(C_1\) は必ず \(1\)(新しい値)です。\(i=2\) では \(C_2 = 1\)(既存)か \(C_2 = 2\)(新しい)のどちらかです。

箱番号列との対応

圧縮番号は箱番号の「出現順リネーム」なので、元の箱番号列を復元する際の自由度は以下の通りです:

  • 新しい圧縮番号 \(k+1\) が登場するとき\(M\) 個の箱のうち、まだ使われていない \(M - k\) 個の箱から1つ選ぶ(\(k\) は現在の種類数)。
  • 既存の圧縮番号 \(d \leq k\) が使われるとき:対応する箱は既に決まっているので、箱の選び方は \(1\) 通り。ただし \(D_i = 0\)(不明)のときは \(k\) 種類のどれでもよいので \(k\) 通り。

DPの状態設計

上記の観察から、ボールを左から順に処理するとき、必要な情報は「現在までに登場した異なる箱番号の種類数 \(k\)」だけです。

\(dp[k]\) = ボール \(i\) まで処理し終えた時点で、ちょうど \(k\) 種類の箱番号が登場している場合の数

アルゴリズム

初期状態は \(dp[0] = 1\)(ボールを1つも処理していない段階では種類数0)。

各ボール \(i\) について \(D_i\) の値に応じて遷移します:

\(D_i = 0\)(不明)の場合: 状態 \(k\) からの遷移は2種類: 1. 既存の値を使う\(k\) 種類のどれかを選ぶ → \(dp'[k] \mathrel{+}= dp[k] \times k\) 2. 新しい値を使う:圧縮番号 \(k+1\) を割り当て、\(M-k\) 個の未使用箱から選ぶ → \(dp'[k+1] \mathrel{+}= dp[k] \times (M - k)\)

\(D_i \neq 0\)(確定値 \(d\))の場合: 状態 \(k\) からの遷移: 1. \(d = k + 1\) のとき(新しい値)\(dp'[k+1] \mathrel{+}= dp[k] \times (M - k)\) 2. \(d \leq k\) のとき(既存の値)\(dp'[k] \mathrel{+}= dp[k] \times 1\) 3. それ以外は不正なので遷移しない。

最終的に \(\sum_k dp[k]\) が答えです。

計算量

  • 時間計算量: \(O(N \times \min(N, M))\)
    各ボールについて状態 \(k\)\(0\) から \(\min(N, M)\) まで。制約で \(N \times \min(N, M) \leq 2.5 \times 10^7\) が保証されています。
  • 空間計算量: \(O(\min(N, M))\)

実装のポイント

  • \(k\) の上限は \(\min(N, M)\) です。\(N\) 個のボールでは最大 \(N\) 種類しか登場せず、箱も \(M\) 個しかないためです。

  • \(D_i = d\) が確定しているとき、\(d > k+1\)\(d = k+1\) だが \(k+1 > K\) などの不正ケースでは遷移が発生せず、自然に \(0\) になります。

  • Pythonでは速度が心配ですが、dp[k] == 0 のスキップにより不要な計算を省けます。

    ソースコード

import sys

def solve():
    input_data = sys.stdin.buffer.read().split()
    N = int(input_data[0])
    M = int(input_data[1])
    D = [int(input_data[i + 2]) for i in range(N)]
    
    MOD = 998244353
    K = min(N, M)
    
    # dp[k] = number of ways, where k = number of distinct compressed values seen so far
    dp = [0] * (K + 2)
    dp[0] = 1
    
    for i in range(N):
        d = D[i]
        new_dp = [0] * (K + 2)
        
        if d == 0:
            # From state k: stay at k with factor k, or go to k+1 with factor (M-k)
            for k in range(K + 1):
                if dp[k] == 0:
                    continue
                val = dp[k]
                # existing value: k choices
                if k > 0:
                    new_dp[k] = (new_dp[k] + val * k) % MOD
                # new value: (M - k) choices, go to k+1
                if k < K and (M - k) > 0:
                    new_dp[k + 1] = (new_dp[k + 1] + val * (M - k)) % MOD
        else:
            # d is fixed
            for k in range(K + 1):
                if dp[k] == 0:
                    continue
                val = dp[k]
                # Case 1: d = k+1, new value
                if d == k + 1 and k + 1 <= K:
                    new_dp[k + 1] = (new_dp[k + 1] + val * (M - k)) % MOD
                # Case 2: d <= k, existing value
                if d <= k:
                    new_dp[k] = (new_dp[k] + val) % MOD
        
        dp = new_dp
    
    ans = 0
    for k in range(K + 2):
        ans = (ans + dp[k]) % MOD
    
    print(ans)

solve()

この解説は claude4.6opus-thinking によって生成されました。

posted:
last update: