【学习笔记】 CF850F Rainbow Balls

我是真不会数学啊

首先可以考虑固定最后剩下的那个颜色。

设$f_n$表示这个颜色的球还剩$n$个，最后所有球都变成这个颜色期望步数。

我们可以观察到，一次操作是好的，当且仅当选择了两个颜色不同的球，并且向左和向右的概率是相等的。

记$s=\sum a_i$。如果这种颜色的球变为了$0$，或者变为了$s$，我们都视为终止。那么我们在已知最终所有球都变为这个颜色的前提下，求出期望$f_{a_i}$以及所有求变成这个颜色的概率$p_{a_i}$（相当于走到$s$而不是$0$的概率），答案是$\sum f_{a_i}{p}_{a_i}$ 。

这里有结论，$p_i=\frac{i}{s}$ 。记${\delta }_i=f_{i+1}-f_i$，不难得到${\delta }_i={\delta }_{i-1}-\frac{s-1}{s-i}$。那么${\delta }_i={\delta }_0-{\sum }_{j=1}^i\frac{s-1}{s-j}$ 。又因为$f_n={\sum }_{i=0}^{n-1}{\delta }_i$，解得${\delta }_0=\frac{(s-1{)}^2}{s}$ 。然后递推即可。

$\text{remark}$ 最后这一步解方程是我始料未及的 甚至花费了很多时间 。并且解出${\delta }_0$意味着我们可以从前往后递推，这样时间复杂度降为$O(\max (a_i))$了。

#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; 
}