「ZJOI2019」开关

考虑进行容斥。记 $F(z)$ 是先假设一直按下去，$n$ 次后恰好所有灯都达到状态的概率的指母函数，则有

$$F(z)=\prod _{i=1}^n\frac{{\mathrm{e}}^{p_{i}z}+(-1{)}^{s_i}{\mathrm{e}}^{-p_{i}z}}{2}$$

注意此处的 $p_i$ 是概率，即原本的 $\frac{p_i}{\sum p_i}$。

而我们想要扣掉第一次达到状态后再次返回到原状态的情况，记 $G(z)$ 是回到原点的概率的指母函数

$$G(z)=\prod _{i=1}^n\frac{{\mathrm{e}}^{p_{i}z}+{\mathrm{e}}^{-p_{i}z}}{2}$$

如果其对应的 OGF $f(z)={\int }_0^{+\infty }F(zt){\mathrm{e}}^{-t}\mathrm{d}t$ 和 $g(z)$，显然 $h=f/g$ 就是第一次到达状态的概率的母函数，$h^{\prime }(1)$ 就是答案。

设 $F(z)=\sum a_w{\mathrm{e}}^{wz}$，则 $f(z)=\sum \frac{a_w}{1-wz}$，因此我们只需要背包出 $F$ 的表示方法，考虑 $h^{\prime }=\frac{f^{\prime }g-g^{\prime }f}{g^2}$，但是 $f,g$ 中含有 $\frac{1}{1-z}$ 项，应当上下乘以 $1-z$，然后计算即可。推导可得 ${{\left(\frac{1-z}{1-wz}\right)}^{\prime }|}_{z=1}=\frac{1}{w-1}$

记 $M=\sum p_i$，复杂度为 $\Theta (nM)$，若用分治 FFT 或者多项式 exp 优化则可以做到 $\Theta (M{\log }^{2}M)$ 或 $\Theta (M\log M)$。

#include 
#include 
#include 
#include 
#include 
#include 

#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 

#define LOG(FMT...) fprintf(stderr, FMT)

using namespace std;

typedef long long ll;
typedef unsigned long long ull;

const int N = 110, M = 50010, P = 998244353;

int n, m;
int s[N], p[N], a[M], b[M];

int norm(int x) { return x >= P ? x - P : x; }

void exGcd(int a, int b, int& x, int& y) {
  if (!b) {
    x = 1;
    y = 0;
    return;
  }
  exGcd(b, a % b, y, x);
  y -= a / b * x;
}

int inv(int a) {
  int x, y;
  exGcd(a, P, x, y);
  return norm(x + P);
}

int calc(int* arr) {
  int ret = 0;
  for (int i = 0; i < m; ++i) {
    int q = norm(2LL * i * inv(m) % P + P - 1);
    ret = (ret + arr[i] * (ll)inv(norm(P + q - 1))) % P;
  }
  return ret;
}

int main() {
  scanf("%d", &n);
  for (int i = 1; i <= n; ++i)
    scanf("%d", &s[i]);
  for (int i = 1; i <= n; ++i)
    scanf("%d", &p[i]);
  a[0] = 1;
  b[0] = 1;
  for (int i = 1; i <= n; ++i) {
    m += p[i];
    if (s[i] == 1) {
      for (int j = m; j >= p[i]; --j)
        a[j] = norm(P + a[j - p[i]] - a[j]);
      for (int j = p[i] - 1; j >= 0; --j)
        a[j] = norm(P - a[j]);
    } else
      for (int j = m; j >= p[i]; --j)
        a[j] = norm(a[j - p[i]] + a[j]);
    for (int j = m; j >= p[i]; --j)
      b[j] = norm(b[j - p[i]] + b[j]);
  }
  int f = calc(a), g = calc(b);
  int ans = norm(f + P - g);
  printf("%d\n", ans);
  return 0;
}