Codeforces 850F. Rainbow Balls

貌似没看到有人用条件概率做这道题的，来整一个

题目链接：F - Rainbow Balls

题目大意：一开始有 $n$ 种不同颜色的球，第 $i$ 种颜色有 $a_i$ 个。每次操作随机先后选取两个球，并将后取出来的球颜色染成前者的颜色。问期望进行多少次操作后所有球的颜色均相同。

符号约定与解释

$P(x)$：事件 $x$ 发生的概率

$P(x,y)$：事件 $x,y$ 同时发生的概率

$P(x|y)$：在事件 $y$ 发生的条件下，事件 $x$ 发生的概率。显然有 $P(x|y)=\frac{P(x,y)}{P(y)}$

$P_{i,j}$：某时刻某颜色的个数为 $i$，且下一时刻个数为 $j$ 的概率。显然只有 $j=i-1,i,i+1$ 是有效的

$N_i$：某颜色的球个数为 $i$

$x_k$：第 $k$ 时刻某颜色的个数

$F_i$：最终局面为所有球都变为第 $i$ 种颜色这一事件，我们也称为第 $i$ 种颜色获胜

$E(N_i|F_j)$：当前颜色为 $j$ 的球的个数为 $i$，在事件 $F_j$ 发生的条件下，期望的操作次数

$P_{i,j}|F_k$：在事件 $F_k$ 发生的条件下，发生颜色个数从 $i$ 到 $j$ 转移的概率

可以得出，最终答案即为 $\sum_{i=1}^{n} E(N_{a_i}|F_i)\times P(F_i)$。

弱化版本

先来考虑这么一个弱化版本的问题：若一开始每种颜色的球都恰好有一个怎么做？

此时，根据对称性质，$P(F_i)$ 固定为 $\frac{1}{n}$，于是此时答案就为 $\sum_{i=1}^{n} E(N_1|F_i)\times \frac{1}{n}=E(N_1|F_1)$ 。

而对任意 $i\in [1,n)$，有：

$$E(N_i|F_1)=1+E(N_{i-1}|F_1)\times P_{i,i-1}|F_1+E(N_{i}|F_1)\times P_{i,i}|F_1+E(N_{i+1}|F_1)\times P_{i,i+1}|F_1 \tag{1}$$

这里需要注意的是，$P_{1,0}|F_1$ 必然是为 $0$ 的，这是因为我们已经制定了 $F_1$ 这一前提条件，所以对应颜色的球个数是不可能为零的。

于是这题的瓶颈就变成了求 $P_{i,j}|F_1$，以求 $P_{i,i-1}|F_1$ 为例，我们首先可以将其改写成 $P(x_{k+1}=i-1|x_k=i,F_1)$，现在我们开始求对应的值。

$$\begin{aligned} P(x_{k+1}=i-1|x_k=i,F_1)& = \frac{P(x_{k+1}=i-1,x_k=i,F_1)}{P(x_k=i,F_1)}\\ & = \frac{P(x_{k+1}=i-1,F_1|x_k=i)\times P(x_k=i)}{P(F_1|x_k=i)\times (x_k=i)}\\ & = \frac{P(F_1|x_{k+1}=i-1,x_{k}=i)\times P(x_{k+1}=i-1|x_{k}=i)}{P(F_1|x_k=i)}\\ & = \frac{\frac{i-1}{n}\times \frac{i(n-i)}{n(n-1)}}{\frac{i}{n}}\\ & = \frac{(n-i)\times (i-1)}{n(n-1)} \end{aligned}$$

过程中化简的思路大致就是把公共项提出来，然后变成计算 $P(F_1|x_k=i)$，即某时刻颜色 $1$ 的个数为 $i$ 时，其获胜的概率。对应值的计算也很简单，因为 $n$ 个颜色不同的球每个球的获胜概率均为 $\frac{1}{n}$，那么此情况下可以看做有其中 $i$ 个颜色被当作颜色 $1$，于是 $P(F_1|x_k=i)=\frac{i}{n}$。

采用类似的计算思路，最终我们可以得出：

$$P_{i,i-1}|F_1=\frac{(n-i)\times (i-1)}{n(n-1)},P_{i,i+1}|F_1=\frac{(n-i)\times (i+1)}{n(n-1)},P_{i,i}|F_1=1-\frac{(n-i)\times 2i}{n(n-1)}$$

那么令 $f_i=E(N_i|F_1)$，将结果代入式 $(1)$，可以得到如下结果：

$f_i=1+\frac{(n-i)\times (i-1)}{n(n-1)}f_{i-1}+\frac{(n-i)\times (i+1)}{n(n-1)}f_{i+1}+(1-\frac{(n-i)\times 2i}{n(n-1)})f_i$

$\frac{(n-i)\times 2i}{n(n-1)}f_i=1+\frac{(n-i)\times (i-1)}{n(n-1)}f_{i-1}+\frac{(n-i)\times (i+1)}{n(n-1)}f_{i+1}$

$2i(n-i)f_i=n(n-1)+(n-i)(i-1)f_{i-1}+(n-i)(i+1)f_{i+1} \tag{2}$

根据初始条件 $f_n=0$，我们最终可以推出 $f_1=(n-1)^2$，具体过程放在最后。

回到本题

令 $m=\sum_{i=1}^{n}a_i$ ，那么可以得出 $P(F_i)=\frac{a_i}{m}$。而根据对称性，$E(N_{a_i}|F_i)=E(N_{a_i}|F_1)$。于是我们就可以得到 $ans=\sum_{i=1}^{n} f_{a_i}\times \frac{a_i}{m}$，而 $f_i$ 的值可以直接根据 $(2)$ 式递推求出，最终我们就能以 $O(\max(a_i)\cdot \log M)$ 的时间复杂度解决此题。

代码实现

#include<bits/stdc++.h>
using namespace std;
#define N 100010
#define LL long long
#define MOD 1000000007
int n,m,x,c[N],f[N],ans;
LL qow(LL x,LL y){return y?(y&1?x*qow(x,y-1)%MOD:qow(x*x%MOD,y/2)):1;}
int main()
{
	scanf("%d",&m);
	for(int i=1;i<=m;i++){
		scanf("%d",&x);
		c[x]++;
		n+=x;
	}
	f[1]=1ll*(n-1)*(n-1)%MOD;
	for(int i=1;i<N-5;i++){
		LL tot=2ll*i*(n-i)%MOD*f[i]%MOD;
		tot=(tot+MOD-1ll*(i-1)*(n-i)%MOD*f[i-1]%MOD)%MOD;
		tot=(tot+MOD-1ll*n*(n-1)%MOD)%MOD;
		tot=tot*qow(i+1,MOD-2)%MOD*qow(n-i,MOD-2)%MOD;
		f[i+1]=tot;
		ans=(ans+qow(n,MOD-2)*i%MOD*f[i]%MOD*c[i])%MOD;
	}
	printf("%d\n",ans);
}

推出 $f_1$ 的过程

$2i(n-i)f_i=n(n-1)+(n-i)(i-1)f_{i-1}+(n-i)(i+1)f_{i+1}$

$2i\cdot f_i=\frac{n(n-1)}{n-i}+(i-1)f_{i-1}+(i+1)f_{i+1}$

令 $t_i=i\cdot f_i$，对 $i\in [1,n)$ 就有 $2t_i=\frac{n(n-1)}{n-i}+t_{i-1}+t_{i+1}$，且 $t_0=t_n=0$

那么设 $t_{i}=a_i\cdot t_{i-1}+b_i$，借由 $t_n=0$ 可以推出 $a_{n-1}=\frac{1}{2},b_{n-1}=\frac{n(n-1)}{2}$。考虑找到 $a_i,b_i$ 的递推式从而推出 $a_2,b_2$ 的值，与 $i=1$ 的情况 $2t_1=n+t_2$ 联立求解

$2t_i=\frac{n(n-1)}{n-i}+t_{i-1}+t_{i+1}=\frac{n(n-1)}{n-i}+t_{i-1}+a_{i+1}\cdot t_i+b_{i+1}$

$a_i=\frac{1}{2-a_{i+1}},b_i=\frac{1}{2-a_{i+1}}(b_{i+1}+\frac{n(n-1)}{n-i})$

手推几项不难瞪眼看出 $a_i=\frac{n-i}{n-i+1},b_i=\frac{n-i}{n-i+1}\cdot n(n-1)$，于是就有

$$\left\{\begin{matrix}2t_1 =n+t_2\\t_2 = \frac{n-2}{n-1}t_1+n(n-2) \end{matrix}\right.$$

最后解得 $f_1=t_1=(n-1)^2$