Solution -「ZJOI 2019」「洛谷 P5326」开关

\(\mathcal{Description}\)

  Link.

  有 \(n\) 个开关，初始时所有开关的状态为 \(0\)。给定开关的目标状态 \(s_1,s_2,\cdots,s_n\)。每次操作中会以正比于 \(p_i\) 的概率拨动开关 \(i\)。求开关达到目标状态的期望操作次数，对 \(998244353\) 取模。

  \(n\le100\)，\(\sum p\le5\times10^4\)。

\(\mathcal{Solution}\)

  不妨令 \(p_i\) 为一次操作拨动 \(i\) 的概率。设 \(F(x)\) 为“\(i\) 次操作后开关是目标状态的概率”的 EGF，\(G(x)\) 为“\(i\) 次操作后回到全零状态的概率”的 EGF。考虑每个开关是否需要被拨动，得到

\[F(x)=\prod_{i=1}^n\frac{e^{p_ix}+(-1)^{s_i}e^{-p_ix}}{2}
\]

代入 \(s_i=0~(i=1,2,\dots,n)\) 即得

\[G(x)=\prod_{i=1}^n\frac{e^{p_ix}+e^{-p_ix}}{2}
\]

  设 \(H(x)\) 为“\(i\) 次操作第一次使开关达到目标状态的概率”的 EGF。\(F\) 和 \(H\) 的区别在于是否接受“多次回到目标状态”，而“回到目标状态”正对应着 \(G\) 的意义，它们可以建立等量关系

\[F=H\cdot G\\
\Rightarrow~~~~H=F\cdot G^{-1}
\]

故欲求期望 \(H'(1)\)，仅需求 \((F\cdot G^{-1})'(1)\)。

  设 \(u_{1-i}=2^n[e^{ix}]F(x)\)，\(v_{1-i}=2^n[e^{ix}]G(x)\)，则有

\[\begin{aligned}
u_{1-i}&=2^n[e^{ix}]F(x)\\
&=[x^{i}]\prod_{j=1}^n(x^{p_j}+(-1)^{s_j}x^{-p_j})\\
&=[x^{i-1}]\prod_{j=1}^n(1+(-1)^{s_j}x^{-2p_j})\\
&=[x^{1-i}]\prod_{j=1}^n(1+(-1)^{s_j}x^{2p_j})
\end{aligned}
\]

同理地，对于 \(v_i\)：

\[\begin{aligned}
v_{1-i}&=2^n[e^{ix}]G(x)\\
&=\cdots\\
&=[x^{1-i}]\prod_{j=1}^n(1+x^{2p_j})
\end{aligned}
\]

在此基础上考虑所求答案：

\[H(x)=\frac{\sum_{i}2^{-n}u_{1-i}e^{ix}}{\sum_i2^{-n}v_{1-i}e^{ix}}
\]

注意 \(e^{ix}=\operatorname{EGF}\lang 1,i,i^2,\cdots\rang\)，将其统一转为 \(\operatorname{OGF}\lang 1,i,i^2,\cdots\rang=\frac{1}{1-ix}\)，此时 \(H\) 的含义变为概率的 OGF。可以得到

\[\begin{aligned}
H(x)&=\frac{\sum_i\frac{u_{1-i}}{1-ix}}{\sum_i\frac{v_{1-i}}{1-ix}}\\
&=\frac{(1-x)\sum_i\frac{u_{1-i}}{1-ix}}{(1-x)\sum_i\frac{v_{1-i}}{1-ix}}\\
&=\frac{u_0+\sum_{i\not=1}u_{1-i}\frac{1-x}{1-ix}}{v_0+\sum_{i\not=1}v_{1-i}\frac{1-x}{1-ix}}
\end{aligned}
\]

记 \(s(x)=u_0+\sum_{i\not=1}u_{1-i}\frac{1-x}{1-ix}\)，\(t(x)=v_0+\sum_{i\not=1}v_{1-i}\frac{1-x}{1-ix}\)。由于有 \(\left(\frac{1-x}{1-ix}\right)'(1)=\frac{1}{1-i}\)，可知

\[s'(1)=\sum_{i\not=1}\frac{u_{1-i}}{1-i}\\
t'(1)=\sum_{i\not=1}\frac{v_{1-i}}{1-i}
\]

而显然又有 \(s(1)=t(1)=u_0=v_0=1\)，则对于 \(H'(1)\)：

\[\begin{aligned}
H'(1)&=\frac{s'(1)t(1)-s(1)t'(1)}{t^2(1)}\\
&=s'(1)-t'(1)\\
&=\sum_{i\not=1}\frac{u_{1-i}-v_{1-i}}{i-1}\\
&=\sum_{i>0}\frac{v_i-u_i}{i}
\end{aligned}
\]

  故仅需求出 \(u_i\) 和 \(v_i\)，有意义的 \(i\) 仅有 \(\mathcal O(\sum p)\)（其中 \(p\) 即输入）个，背包一下，即可 \(\mathcal O(nm)\) 求解。

\(\mathcal{Code}\)

/* Clearink */

#include <cstdio>

#define rep( i, l, r ) for ( int i = l, rpbound##i = r; i <= rpbound##i; ++i )
#define per( i, r, l ) for ( int i = r, rpbound##i = l; i >= rpbound##i; --i )

const int MAXN = 100, MAXS = 5e4, MOD = 998244353;
int n, s[MAXN + 5], p[MAXN + 5], u[MAXS + 5], v[MAXS + 5];

inline int mul( const long long a, const int b ) { return a * b % MOD; }
inline int sub( int a, const int b ) { return ( a -= b ) < 0 ? a + MOD : a; }
inline int add( int a, const int b ) { return ( a += b ) < MOD ? a : a - MOD; }
inline int mpow( int a, int b ) {
	int ret = 1;
	for ( ; b; a = mul( a, a ), b >>= 1 ) ret = mul( ret, b & 1 ? a : 1 );
	return ret;
}

int main() {
	scanf( "%d", &n );
	rep ( i, 1, n ) scanf( "%d", &s[i] );
	int sp = 0;
	rep ( i, 1, n ) scanf( "%d", &p[i] ), sp += p[i];
	u[0] = v[0] = 1;
	rep ( i, 1, n ) per ( j, sp, p[i] ) {
		u[j] = ( s[i] ? sub : add )( u[j], u[j - p[i]] );
		v[j] = add( v[j], v[j - p[i]] );
	}
	// rep ( i, 0, sp ) printf( "%d%c", u[i], i ^ sp ? ' ' : '\n' );
	// rep ( i, 0, sp ) printf( "%d%c", v[i], i ^ sp ? ' ' : '\n' );
	int ans = 0;
	rep ( i, 1, sp ) {
		ans = add( ans, mul( mpow( i << 1, MOD - 2 ), sub( v[i], u[i] ) ) );
	}
	printf( "%d\n", mul( ans, sp ) );
	return 0;
}