【学习笔记】 CF850F Rainbow Balls

前置知识：条件期望。（应该是对的叭？）

记$s=\sum a_i$。考虑一个颜色的球有$i$个，所有球变成它颜色的概率，答案是$p_i=\frac{i}{s}$。这是因为，可以将所有球看成两两不同，最终变成每个球的颜色的概率是相同的。

可以考虑固定最后剩下的那个颜色。设$f_n$表示这个颜色的球还剩$n$个，最后所有球都变成这个颜色期望步数。答案就是$\sum f_{a_i}$。

我们可以观察到，一次操作是好的，当且仅当选择了两个颜色不同的球，并且向左和向右的概率是相等的（指下标）。因此递推式如下：

$$f_i=\frac{\left(\genfrac{}{}{0ex}{}{s}{2}\right)}{i(s-i)}\times \frac{i}{s}+\frac{1}{2}(f_{i-1}+f_{i+1})$$

注意计算条件期望的时候，每一步转移的代价要乘上对应的能到达目标状态的概率（感性理解）。可以认为$0$和$s$都是终止状态，我们要乘上的就是到达$s$的概率。

记${\delta }_i=f_{i+1}-f_i$，将上式变形可以得到：

$${\delta }_i={\delta }_{i-1}-\frac{s-1}{s-i}={\delta }_0-\sum _{j=1}^i\frac{s-1}{s-j}$$

根据${\sum }_{i=0}^{s-1}{\delta }_i=0$，事实上已经可以把${\delta }_0$解出来了。解出来${\delta }_0=\frac{(s-1{)}^2}{s}$，从前往后递推即可。

复杂度$O(\max (a_i))$。

upd on 2024/2/6：没有看懂自己写的题解，警钟长鸣。

#include<bits/stdc++.h>
#define fi first
#define se second
#define ll long long
#define pb push_back
#define db double
#define inf 0x3f3f3f3f3f3f3f3f
using namespace std;
const int mod=1e9+7;
int n,m,a[2505];
ll s,f[100005],g[100005],res;
ll fpow(ll x,ll y){
	ll z(1);
	for(;y;y>>=1){
		if(y&1)z=z*x%mod;
		x=x*x%mod;
	}return z;
}
int main(){
	ios::sync_with_stdio(false);
	cin.tie(0),cout.tie(0);
	cin>>n;for(int i=1;i<=n;i++)cin>>a[i],m=max(m,a[i]),s+=a[i];
	g[0]=(s-1)*(s-1)%mod*fpow(s,mod-2)%mod;
	for(int i=1;i<m;i++)g[i]=(g[i-1]-(s-1)*fpow(s-i,mod-2)%mod)%mod;
	for(int i=1;i<=m;i++)f[i]=(f[i-1]+g[i-1])%mod;
	for(int i=1;i<=n;i++)res=(res+f[a[i]])%mod;
	cout<<(res+mod)%mod; 
}