[CF850F] Rainbow Balls

题目链接

题目描述

给定 $n$ 种颜色的球，每种球有 $a_i$ 个，对这些球执行以下操作：

• 有顺序地任意取两个球，将第二个球涂上第一个球的颜色，重复该操作至所有球颜色相同。

求期望操作次数，对 $10^9+7$ 取模。
数据范围：$n\le 2500$，$1\le a_i\le 10^5$。

Solution

• 设 $f_i$ 表示当前有 $i$ 个球，将所有球变为该颜色的期望次数，$s$ 表示球的总数，$p$ 表示当前取出的两个球第一个与最终颜色相同，第二个与最终颜色不同的概率。
  根据题意，有 $$p=\frac{i\times (s-i)}{s\times (s-1)}$$$$f_i=p\times f_{i-1}+p\times f_{i+1}+(1-2p)\times f_i+v$$其中 $v$ 为这一步操作对最终答案的贡献，也就是该颜色成为最终颜色的概率。
• 设 $g_i$ 表示当前颜色成为最终颜色的概率，有 $g_0=0$，$g_1=1$，且 $g_i=p\times g_{i-1}+p\times g_{i+1}+(1-2p)\times g_i$。
  转化可得 $g_i-g_{i-1}=g_{i+1}-g_i$，等差数列求和可得 $g_i=\frac{i}{s}$，即 $v=\frac{i}{s}$。
• 那么现在有$$f_i=p\times f_{i-1}+p\times f_{i+1}+(1-2p)\times f_i+\frac{i}{s}$$转化后有$$f_i-f_{i+1}=f_{i-1}-f_i+\frac{s-1}{s-i}$$
• 考虑求 $f_1$ 和 $f_2$ 后线性递推，由于 $f_0$ 不存在，代入后有 $f_2=2f_1-1$。
  $$\begin{array}{rl}{f}_1& =f_1-f_s\\ & =\sum _{i=2}^s{f}_{i-1}-f_i\\ & =(s-1)\times (f_1-f_2)+\sum _{i=2}^{s-1}\frac{s-1}{s-i}\times (s-i)\\ & =(s-1)\times (1-f_1)+(s-1)\times (s-2)\end{array}$$
• 所以$$f_1=\frac{(s-1{)}^2}{s}$$线性递推即可求解，答案为 ${\sum }_{i=1}^n{f}_{a_i}$。

Code

#include<cstdio>
using namespace std;
const int maxn=100010,MLY=1000000007;
int f[maxn],n,a[maxn],s,ans;
inline int power(int a,int b){
	int ans=1;
	while(b){
		if(b&1)ans=1ll*ans*a%MLY;
		a=1ll*a*a%MLY;
		b>>=1;
	}
	return ans;
}
int main(){
	scanf("%d",&n);
	for(int i=1;i<=n;++i)scanf("%d",&a[i]),s+=a[i];
	f[1]=(s-1ll)*(s-1)%MLY*power(s,MLY-2)%MLY;
	f[2]=(2ll*f[1]-1+MLY)%MLY;
	for(int i=2;i<100000;++i)
		f[i+1]=((2ll*f[i]-f[i-1]-(s-1ll)*power(s-i,MLY-2))%MLY+MLY)%MLY;
	for(int i=1;i<=n;++i)ans=(ans+f[a[i]])%MLY;
	printf("%d",ans);
	return 0;
}