[AtCoder Educational DP Contest] J - Sushi（期望dp）

problem

luogu

现有$N(1\le N\le 300)$ 个盘子，编号为$1,2,3,\dots ,N$。

第 $i$个盘中放有 $a_i(1\le a_i\le 3)$个寿司。

接下来每次执行以下操作，直至吃完所有的寿司。

从第 $1,2,3,\dots ,N$ 个盘子中任选一个盘子，吃掉其中的一个寿司。若没有寿司则不吃。

若将所有寿司吃完，请问此时操作次数的数学期望是多少？

solution

最直接地设 $f(a_1,a_2,a_3,...,a_n):$ 第 $i$ 盘还剩 $a_i$ 个寿司的期望次数。

那么枚举随机到的盘子，有方程：$f(a_1,a_2,...,a_n)=1+{\sum }_{i=1}^n\frac{1}{n}f(a_1,a_2,...,\max (a_i-1,0),...,a_n)$。

显然这个等式不能构成转移方程，因为存在原地打转（第 $i$ 盘寿司为空时，就变成了自身转移到自身，状态不变）的情况。

由于随机均匀分布，盘子的位置是不重要的，事实上我们真正关注的只有盘子中剩余寿司的数量。

而 $a_i$ 寿司数量又只有四种取值 $0/1/2/3$。

不妨重新设 $f(o,i,j,k):$ 当前还剩下 $o/i/j/k$ 个盘子中有 $0/1/2/3$ 个寿司。

则有转移：
$$f(o,i,j,k)=1+\frac{o}{n}f(o,i,j,k)+\frac{i}{n}f(o+1,i-1,j,k)+\frac{j}{n}f(o,i+1,j-1,k)+\frac{k}{n}f(o,i,j+1,k-1)$$

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

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

这样就不存在状态相同的死循环转移了。

但是现在是 $O(n^4)$ 的，需要进一步优化。

不难发现，盘子数量是固定不变的，即 $o+i+j+k=n$，所以当我们知道了其中任意三个数，就能推出剩下一个数。

设 $f(i,j,k):$ 当前还剩下 $i/j/k$ 个盘子中有 $1/2/3$ 个寿司。
$$f(i,j,k)=\frac{n}{i+j+k}+\frac{i}{i+j+k}f(i-1,j,k)+\frac{j}{i+j+k}f(i+1,j-1,k)+\frac{k}{i+j+k}f(i,j+1,k-1)$$
最后还要注意循环枚举的细节：

• $k$ 只用了 $k-1$，当在 $(i,j)$ 时会问到 $j+1$，所以 $k$ 要在 $j$ 循环的外层。
• 同理 $j$ 会用到同 $k$ 下的 $j+1$，但此时是 $i-1$，所以 $j$ 循环要在 $i$ 的外层。
• 综上我们确定唯一的循环顺序是 $k,j,i$。

时间复杂度： $O(n^3)$。

code

#include <bits/stdc++.h>
using namespace std;
#define maxn 305
double f[maxn][maxn][maxn];
int a[5];
int n;
int main() {
	scanf( "%d", &n );
	for( int i = 1, x;i <= n;i ++ ) scanf( "%d", &x ), a[x] ++;
	for( int k = 0;k <= n;k ++ )
	for( int j = 0;j <= n;j ++ )
	for( int i = 0;i <= n;i ++ )
		if( i or j or k ) {
			if( i ) f[i][j][k] += f[i - 1][j][k] * i / (i + j + k);
			if( j ) f[i][j][k] += f[i + 1][j - 1][k] * j / (i + j + k);
			if( k ) f[i][j][k] += f[i][j + 1][k - 1] * k / (i + j + k);
			f[i][j][k] += n * 1.0 / (i + j + k);
		}	
	printf( "%.10f\n", f[a[1]][a[2]][a[3]] );
	return 0;
}