AtCoder 🟢 AT_dp_j Sushi

題目的難度顏色使用 Luogu 上的分級，由簡單到困難分別為 🔴🟠🟡🟢🔵🟣⚫。

🔗 🟢 AT_dp_j Sushi

Problem Statement

題目簡述

有 $N$ 個盤子，每個盤子上最初放有 $a_i \in \{1, 2, 3\}$ 個壽司。
每次操作擲一枚 $1 \sim N$ 的公平骰子，若骰到的盤子上有壽司，則吃掉其中一個；若沒有則不進行任何動作。
求將所有壽司吃完所需的期望操作次數。

Constraints

約束條件

• $1 \le N \le 300$
• $1 \le a_i \le 3$

---

思路：期望值 DP

這是一道典型的期望值 DP 問題。

狀態壓縮的關鍵觀察

由於每個盤子的壽司數量只有 $1, 2, 3$ 三種，我們不需要記錄每個盤子的具體狀態，只需記錄各類盤子的數量即可。

狀態定義

令 $f(i, j, k)$ 表示當前盤面上有 $i$ 個剩 1 個壽司的盤子、$j$ 個剩 2 個壽司的盤子、$k$ 個剩 3 個壽司的盤子時，吃完所有壽司所需的期望操作次數。

• 目標：求 $f(\text{cnt}_1, \text{cnt}_2, \text{cnt}_3)$，其中 $\text{cnt}_x$ 為初始時有 $x$ 個壽司的盤子數量。
• 邊界：$f(0, 0, 0) = 0$（沒有壽司時不需操作）。

轉移方程

在狀態 $(i, j, k)$ 下擲一次骰子，根據結果分為四種情況：

結果	機率	後續狀態
骰到空盤子	$\dfrac{N - i - j - k}{N}$	狀態不變
骰到剩 1 個壽司的盤子	$\dfrac{i}{N}$	$(i-1, j, k)$
骰到剩 2 個壽司的盤子	$\dfrac{j}{N}$	$(i+1, j-1, k)$
骰到剩 3 個壽司的盤子	$\dfrac{k}{N}$	$(i, j+1, k-1)$

根據期望的定義，每次操作消耗 1 步：

$$f(i,j,k) = 1 + \frac{N-i-j-k}{N} f(i,j,k) + \frac{i}{N} f(i-1,j,k) + \frac{j}{N} f(i+1,j-1,k) + \frac{k}{N} f(i,j+1,k-1)$$

移項化簡

將含 $f(i,j,k)$ 的項移至左邊：

$$f(i,j,k) \cdot \frac{i+j+k}{N} = 1 + \frac{i}{N} f(i-1,j,k) + \frac{j}{N} f(i+1,j-1,k) + \frac{k}{N} f(i,j+1,k-1)$$

兩邊同乘 $N$ 再除以 $(i+j+k)$，得到最終遞推式：

$$\boxed{f(i,j,k) = \frac{N + i \cdot f(i-1,j,k) + j \cdot f(i+1,j-1,k) + k \cdot f(i,j+1,k-1)}{i+j+k}}$$

實現細節

迭代順序

計算 $f(i, j, k)$ 時需要用到三個子問題：

依賴項	與當前狀態的關係	如何確保已計算
$f(i-1, j, k)$	$i$ 減少，$j, k$ 不變	$i$ 從小到大迭代
$f(i+1, j-1, k)$	$j$ 減少（$i$ 增加無妨）	$j$ 從小到大迭代，同層 $i$ 全部算完
$f(i, j+1, k-1)$	$k$ 減少（$j$ 增加無妨）	$k$ 從小到大迭代，同層 $j$ 全部算完

關鍵觀察：雖然後兩項中 $i$ 或 $j$ 會增加，但對應的 $j$ 或 $k$ 會減少。由於我們按 $k \to j \to i$ 的順序從小到大迭代，當計算 $(i, j, k)$ 時：

• 所有 $k' < k$ 的狀態都已算完 → $f(i, j+1, k-1)$ ✓
• 在同一 $k$ 下，所有 $j' < j$ 的狀態都已算完 → $f(i+1, j-1, k)$ ✓
• 在同一 $(j, k)$ 下，所有 $i' < i$ 的狀態都已算完 → $f(i-1, j, k)$ ✓

迭代範圍優化

實現中使用 k in range(cnt[3] + 1) 而非 range(n + 1)，這是一種空間剪枝優化：

• 從初始狀態 $(cnt_1, cnt_2, cnt_3)$ 出發，$k$ 最多為 $cnt_3$（不可能憑空產生有 3 個壽司的盤子）
  同理，$j \le cnt_2 + cnt_3$，$i \le cnt_1 + cnt_2 + cnt_3$
• 此外，當有 $k$ 個盤子剩 3 個壽司時，不可能有 $j > n - k$ 個盤子剩 2 個壽司，因為 $j + k \le n$（總盤子數不變）
  同理，不可能有 $i > n - j - k$ 個盤子剩 1 個壽司。

這樣只計算可達狀態，避免浪費計算資源

複雜度分析

• 時間複雜度：$\mathcal{O}(N^3)$
  • 狀態數為滿足 $i + j + k \le N$ 的非負整數三元組數量。
  • 令 $\ell = N - i - j - k$（空盤子數），則等價於求 $i + j + k + \ell = N$ 的非負整數解數，即重複組合數 $H^{4}_{N} = \binom{N+3}{3} \approx \frac{N^3}{6}$。
  • 每個狀態的轉移為 $O(1)$。
• 空間複雜度：$\mathcal{O}(N^3)$，用於存儲 DP 表。

---

Code

def solve():
    n = int(input())
    A = list(map(int, input().split()))
    assert len(A) == n

    cnt = [0] * 4
    for a in A:
        cnt[a] += 1

    f = [[[0] * (n + 1) for _ in range(n + 1)] for _ in range(n + 1)]
    # for k in range(n + 1):
    #     for j in range(n - k + 1):
    #         for i in range(n - k - j + 1):
    s1 = cnt[1] + cnt[2] + cnt[3]
    s2 = cnt[2] + cnt[3]
    for k in range(cnt[3] + 1):
        for j in range(min(s2, n - k) + 1):
            for i in range(min(s1, n - k - j) + 1):
                if i == 0 and j == 0 and k == 0:
                    continue
                v = n
                if i > 0:
                    v += i * f[i-1][j][k]
                if j > 0:
                    v += j * f[i+1][j-1][k]
                if k > 0:
                    v += k * f[i][j+1][k-1]
                f[i][j][k] = v / (i + j + k)
    print(f[cnt[1]][cnt[2]][cnt[3]])

if __name__ == "__main__":
    solve()