【CF850F】Rainbow Balls（期望DP）

点此看题目

• 有\(n\)种颜色的球，第\(i\)种颜色的球有\(a_i\)个。
• 一次操作随机选取两个球\(x,y\)，使\(y\)的颜色变得与\(x\)相同。
• 求期望多少次操作之后，所有球颜色相同。
• \(n\le2500,a_i\le10^5\)

初步\(DP\)思路

记\(S=\sum_{i=1}^na_i\)。

考虑枚举每一种颜色，计算所有球都变成这种颜色的概率。

那么有用的信息就仅有场上有多少个目标颜色的球。

因此设\(f_i\)表示有\(i\)个目标颜色的球时使得所有球变成目标颜色的期望步数，\(g_i\)表示有\(i\)个目标颜色的球时使得所有球变成目标颜色的概率。

\(g_i\)的求解

先考虑\(g_i\)，容易列出方程：

\[g_i=\frac{i\times(S-i)}{S\times(S-1)}\times (g_{i-1}+g_{i+1})+(1-2\times \frac{i\times(S-i)}{S\times(S-1)})\times g_i \]

化简一下得到：

\[g_i=\frac{g_{i-1}+g_{i+1}}{2} \]

又因为有边界条件\(g_0=0,g_S=1\)，得出\(g_i=\frac iS\)。

\(f_i\)的求解

注意到我们要求所有球都变成目标颜色，也就是说没有变成目标颜色的情况是不予考虑的，因此这里一步的贡献并非常规\(DP\)中的\(1\)，而是能变成目标颜色的概率\(g_i\)：

\[f_i=\frac{i\times(S-i)}{S\times(S-1)}\times (f_{i-1}+f_{i+1})+(1-2\times \frac{i\times(S-i)}{S\times(S-1)})\times f_i+\frac iS \]

化简一下得到：

\[2f_i=f_{i-1}+f_{i+1}+\frac{S-1}{S-i} \]

通过移项也可以发现：

\[f_i-f_{i+1}=f_{i-1}-f_i+\frac{S-1}{S-i} \]

我们的已知条件是\(f_0=f_S=0\)，由于\(f_1\)和\(f_0\)已经存在关系，因此我们想要通过在\(f_1\)和\(f_S\)之间建立联系列出新的关系式：

\[\begin{aligned} f_1-f_S&=\sum_{i=1}^{S-1}(f_i-f_{i+1})\\ &=\sum_{i=1}^{S-1}(f_0-f_1+\sum_{j=1}^i\frac{S-1}{S-j})\\ &=(S-1)\times(f_0-f_1)+\sum_{i=1}^{S-1}\frac{S-1}{S-i}\times(S-i)\\ &=(S-1)\times(f_0-f_1)+(S-1)^2 \end{aligned} \]

最终得到：

\[f_1=\frac{(S-1)^2}S \]

有了\(f_0,f_1\)，就可以利用前面的方程递推出\(f\)的每一项了。

最终答案就是\(\sum_{i=1}^nf_{a_i}\)。

代码：\(O(VlogV)\)

#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define Rg register
#define RI Rg int
#define Cn const
#define CI Cn int&
#define I inline
#define W while
#define N 2500
#define V 100000
#define X 1000000007
using namespace std;
int n,a[N+5],f[V+5];
I int QP(RI x,RI y) {RI t=1;W(y) y&1&&(t=1LL*t*x%X),x=1LL*x*x%X,y>>=1;return t;}
int main()
{
	RI i,s=0,Mx=0;for(scanf("%d",&n),i=1;i<=n;++i) scanf("%d",a+i),s+=a[i],Mx=max(Mx,a[i]);
	for(f[1]=1LL*(s-1)*(s-1)%X*QP(s,X-2)%X,i=2;i<=Mx;++i) f[i]=(2LL*f[i-1]-f[i-2]-1LL*(s-1)*QP(s-(i-1),X-2)%X+2*X)%X;//求出f[1]，然后递推
	RI t=0;for(i=1;i<=n;++i) t=(t+f[a[i]])%X;return printf("%d\n",t),0;//枚举目标颜色
}