E - 圧縮番号列の復元 / Restoration of Compressed Number Sequence Editorial by admin
Gemini 3.0 Flash (Thinking)概要
この問題は、ボールを箱に入れていく際の「箱の現れ方のパターン(圧縮番号列)」が与えられた条件を満たすような、具体的な「箱番号の割り当て方」の総数を求める問題です。
各ステップで、新しい箱を使うか、すでに使った箱を再利用するかを選択していきます。状態として「これまでに何種類の箱を使ったか」を保持する動的計画法(DP)を用いることで、効率的に解くことができます。
考察
1. 圧縮番号列のルールの理解
圧縮番号列 \(C_i\) は、以下のルールで決まります。 - ボール \(i\) を新しい箱(まだ一度も使っていない箱)に入れる場合: \(C_i\) は「これまでに使った箱の種類数 + 1」となり、この箱にその番号が割り当てられます。 このとき、残っている \(M - (\text{既に使用した種類数})\) 個の箱から \(1\) つ選ぶことになります。 - ボール \(i\) を既に使用した箱に入れる場合: \(C_i\) はその箱に以前割り当てられた番号になります。 特定の圧縮番号 \(d\) に対応する箱は一意に決まっているため、選べる箱は \(1\) 通りしかありません。
2. DPの状態定義
ボールを \(i\) 個目まで入れたとき、それまでに使った箱の種類数を \(k\) とします。 このとき、次に選べる行動は以下の通りです。 - 新しい箱を使う場合:種類数は \(k+1\) に増える。箱の選び方は \((M-k)\) 通り。 - 既に使用した \(k\) 種類のうち、番号 \(d\) の箱を使う場合:種類数は \(k\) のまま。箱の選び方は \(1\) 通り。
したがって、以下のDPを考えます。 - \(dp[k] = \)(これまでに \(k\) 種類の箱を使用しているような箱番号列の総数)
3. 条件 \(D_i\) に基づく遷移
各ボール \(i = 1, \dots, N\) について、与えられた \(D_i\) の値によって遷移を更新します。
ケース1:\(D_i = 0\)(不明な場合)
\(C_i\) は何でも良いため、「既存の \(k\) 種類のどれかを使う」か「新しい箱を使う」かの両方を考慮します。 - 新しい種類数 \(k\) になるのは: - 前の状態が \(k\) 種類で、既存の \(k\) 個のどれかを使う(\(k\) 通り) - 前の状態が \(k-1\) 種類で、新しい箱を使う(\(M-(k-1)\) 通り) - 更新式:\(dp_{new}[k] = dp[k] \times k + dp[k-1] \times (M - k + 1)\)
ケース2:\(D_i = d > 0\)(確定している場合)
\(C_i\) が \(d\) であるためには、以下のいずれかである必要があります。 1. \(d \leq (\text{現在の種類数 } k)\) のとき: すでに番号 \(d\) が割り当てられた箱を使うしかないため、箱の選び方は 1通り です。種類数は \(k\) のまま変わりません。 2. \(d = (\text{前の種類数 } k-1) + 1\) のとき: 新しい箱を使い、それに番号 \(d\) が割り当てられるケースです。箱の選び方は \((M-(d-1))\) 通り です。種類数は \(d\) になります。
これらをまとめると、 - \(k < d\) の状態は、番号 \(d\) が出現したという条件に矛盾するため、すべて \(0\) になります。 - \(k = d\) のとき:\(dp_{new}[d] = dp[d] \times 1 + dp[d-1] \times (M - d + 1)\) - \(k > d\) のとき:\(dp_{new}[k] = dp[k] \times 1\)(既存の箱 \(d\) を選ぶ \(1\) 通りのみ)
アルゴリズム
- DPテーブル
dpをサイズ \(\min(N, M) + 1\) で初期化し、dp[0] = 1とする。 - 各ボール \(i = 1 \dots N\) について順に処理する。
- \(D_i = 0\) の場合:
- \(k\) が大きい方から順に \(dp[k] = (dp[k] \cdot k + dp[k-1] \cdot (M-k+1))\) と更新する。
- \(D_i = d > 0\) の場合:
- \(d > M\) または \(d > (\text{これまでの最大種類数} + 1)\) なら、条件を満たさないので答えは \(0\)。
- \(dp[d] = (dp[d] + dp[d-1] \cdot (M-d+1))\) と更新する。
- \(k < d\) の範囲の
dp[k]をすべて \(0\) にする。
- \(D_i = 0\) の場合:
- 最終的な
dpテーブルの全要素の和が答え。
計算量
- 時間計算量: \(O(N \times \min(N, M))\)
- 制約 \(N \times \min(N, M) \leq 2.5 \times 10^7\) により、この計算量で実行時間制限内に収まります。
- 空間計算量: \(O(\min(N, M))\)
- DPテーブルを 1 次元で保持することで節約しています。
実装のポイント
更新順序: \(D_i = 0\) のときの更新では、同じステップの更新結果を再利用しないよう、インデックス \(k\) を大きい方からループさせるか、新しい配列を用意する必要があります。
高速化: Pythonの場合、内側のループや配列操作を工夫することで実行速度を向上させられます。このコードでは
dp[:d] = [0] * dのようなスライス操作を利用して高速化を図っています。上限の制限: 使える箱の種類数は \(M\) 個までであり、またボールの数 \(N\) を超えることもありません。よって \(k\) の上限は \(\min(N, M)\) となります。”`
ソースコード
import sys
def solve():
# Use fast I/O to read the entire input at once
input_data = sys.stdin.read().split()
if not input_data:
return
# Ball count N and box count M
N = int(input_data[0])
M = int(input_data[1])
MOD = 998244353
# The number of distinct boxes used cannot exceed min(N, M).
# This limits the size of our DP table.
limit = min(N, M)
dp = [0] * (limit + 1)
# Initial state: 0 balls processed, 0 distinct boxes used.
dp[0] = 1
# Current maximum number of distinct boxes that could have been used.
curr_max_k = 0
# Precompute (M - k + 1) % MOD for each possible k.
# This is used when a box with rank k is encountered for the first time.
M_minus_k_plus_1 = [(M - k + 1) % MOD for k in range(limit + 1)]
# Local variables for faster access inside the loop
_MOD = MOD
_M = M
_limit = limit
# Process balls one by one
for i in range(N):
d = int(input_data[i + 2])
if d == 0:
# Case where the compressed box number C_i is unknown.
# We must consider both a new box and any of the existing boxes.
# If we add a new box, its compressed number will be curr_max_k + 1.
# This is only possible if curr_max_k < M and curr_max_k < N.
if curr_max_k < _limit:
dp[curr_max_k + 1] = (dp[curr_max_k] * (_M - curr_max_k)) % _MOD
# If we use an existing box, C_i can be any rank from 1 to curr_max_k.
# dp[k] = dp[k] * k (choosing one of k existing boxes)
# + dp[k-1] * (M - k + 1) (choosing a new box to be the k-th rank).
# We iterate backwards to use the values from the previous ball.
for k in range(curr_max_k, 0, -1):
dp[k] = (dp[k] * k + dp[k-1] * M_minus_k_plus_1[k]) % _MOD
# dp[0] represents 0 boxes used, which is impossible after the first ball.
dp[0] = 0
# Increment the maximum possible number of distinct boxes used.
if curr_max_k < _limit:
curr_max_k += 1
else:
# Case where the compressed box number C_i is given as d.
# C_i must follow the rules: d <= M and d <= curr_max_k + 1.
if d > _M or d > curr_max_k + 1:
print(0)
return
# For any valid state with k distinct boxes used:
# 1. If d == k + 1, ball i is a new box (M - k choices).
# 2. If d <= k, ball i is the box previously assigned rank d (1 choice).
# All other outcomes for ball i are invalid under the condition C_i = d.
# Update dp[d] using the previous states dp[d] and dp[d-1].
new_dp_d = (dp[d] + dp[d-1] * M_minus_k_plus_1[d]) % _MOD
# States where k < d are now impossible because d ranks must have appeared.
dp[:d] = [0] * d
# Assign the updated value for dp[d].
dp[d] = new_dp_d
# Update curr_max_k if the given rank d exceeds the previous max.
if d > curr_max_k:
curr_max_k = d
# The answer is the sum of ways for all possible counts of distinct boxes.
print(sum(dp) % _MOD)
if __name__ == '__main__':
solve()
この解説は gemini-3-flash-thinking によって生成されました。
posted:
last update: