「ZJOI2019」开关
考虑进行容斥。记 F ( z ) F(z) F(z) 是先假设一直按下去， n n n 次后恰好所有灯都达到状态的概率的指母函数，则有
F ( z ) = ∏ i = 1 n e p i z + ( − 1 ) s i e − p i z 2 F(z) = \prod_{i=1}^n \frac{\mathrm e^{p_iz} + (-1)^{s_i} \mathrm e^{-p_iz}}2 F(z)=i=1∏n2epiz+(−1)sie−piz
注意此处的 p i p_i pi 是概率，即原本的 p i ∑ p i \frac{p_i}{\sum p_i} ∑pipi。
而我们想要扣掉第一次达到状态后再次返回到原状态的情况，记 G ( z ) G(z) G(z) 是回到原点的概率的指母函数
G ( z ) = ∏ i = 1 n e p i z + e − p i z 2 G(z) = \prod_{i=1}^n \frac{\mathrm e^{p_iz} + \mathrm e^{-p_iz}}2 G(z)=i=1∏n2epiz+e−piz
如果其对应的 OGF f ( z ) = ∫ 0 + ∞ F ( z t ) e − t d t f(z) = \int_0^{+\infty} F(zt)\mathrm e^{-t} \mathrm d t f(z)=∫0+∞F(zt)e−tdt 和 g ( z ) g(z) g(z)，显然 h = f / g h=f/g h=f/g 就是第一次到达状态的概率的母函数， h ′ ( 1 ) h'(1) h′(1) 就是答案。
设 F ( z ) = ∑ a w e w z F(z) = \sum a_w \mathrm e^{wz} F(z)=∑awewz，则 f ( z ) = ∑ a w 1 − w z f(z) = \sum \frac{a_w}{1- wz} f(z)=∑1−wzaw，因此我们只需要背包出 F F F 的表示方法，考虑 h ′ = f ′ g − g ′ f g 2 h'=\frac{f'g-g'f}{g^2} h′=g2f′g−g′f，但是 f , g f, g f,g 中含有 1 1 − z \frac 1{1-z} 1−z1 项，应当上下乘以 1 − z 1-z 1−z，然后计算即可。推导可得 ( 1 − z 1 − w z ) ′ ∣ z = 1 = 1 w − 1 \left.\left( \frac{1-z}{1-wz} \right)'\right|_{z=1} = \frac1{w-1} (1−wz1−z)′∣∣∣z=1=w−11
记 M = ∑ p i M = \sum p_i M=∑pi，复杂度为 Θ ( n M ) \Theta(nM) Θ(nM)，若用分治 FFT 或者多项式 exp 优化则可以做到 Θ ( M log 2 M ) \Theta(M\log^2 M) Θ(Mlog2M) 或 Θ ( M log M ) \Theta(M\log M) Θ(MlogM)。
#include