[ZJOI2019] 开关

Posted 2020-12-01 vexoben

令 (P = sum p_i)

设 (F) 为走 (x) 步，刚好达到终止状态的概率的 EGF； (G) 为走 (x) 步，刚好走回初始状态的概率的 EGF，则有：
[ F = frac{1}{2}prod_i (e^{frac{p_i}{P}x} + (-1) ^ {s_i} e ^{-frac{p_i}{P}x}) \ G = frac{1}{2} prod_i (e^{frac{p_i}{P}x} + e ^{-frac{p_i}{P}x}) ]
设两者的 OGF 分别为 (f) 和 (g), 再令(h_n) 为走 (n) 步，第一次到达终止状态的概率。那么答案就是(sum_{i ≥ 0} h(i) i = h'(1))。容斥一下得到：
[ h_n = f_n - sum_{i = 0} ^ {n - 1} h_i g_{n-i} ]
于是得到 (h * g = f, h = frac{f}{g})

根据除法的求导法则，得到：
[ h'(1) = frac{f'(1)g(1) - f(1)g'(1)}{g(1) ^ 2} ]
现在只需求得 (f(1),f'(1),g(1),g'(1))

问题在于，我们第一个式子可以在 (O(nP)) 的复杂度内求出 (F, G)， 怎么用 OGF 推出 EGF 呢？

设 (F = sum_{i = -P} ^ {P} a_i e^{frac{i}{P} x} = sum_{i = -P} ^ P a_i sum_{j≥0}frac{(frac{ix}{P})^j}{j!})，就得到：
[ f = sum_{i = -P} ^ P frac{a_i}{1 - frac{i}{P} x} ]
因为当 (x = 1) 时，(f,g) 分母的部分可能会等于 (0). 不妨将两个 OGF 同时乘以 $ prod_{i = -P} ^ P (1 + frac{i}{P} x)$ (注意符号)

我们只需要计算这样的 (f(1), f'(1)):
[ f = sum_{i = -P} ^ P b_i prod_{-P ≤j ≤ P, j ≠ i} (1 + frac{j}{P} x) ]
其中 (b_i = a_{-i}), 这是因为我们刚才乘那个多项式的时候处理了一下符号。

先计算 (f(1)). 当 (x = 1), 显然只有当 (i = -P) 的时候 (prod) 里面才不为 (0)

计算 (f'(1)) 时可以使用这样一个引理：
[ (prod_i(1 + a_ix))' = sum_i a_i prod_{j≠i} (1 + a_j x) ]
证明可以暴力展开：
[ egin{split} (prod_i(1 + a_ix))' &= (e^{sum_i ln(1 + a_ix)})' \ &= (e^{sum_i ln(1 + a_ix)}) (sum_i ln'(1 + a_i x)) \ &= (prod_i (1 + a_ix))(sum_i frac{a_i}{1 + a_ix}) \ &= sum_i a_i prod_{j≠i} (1 + a_j x) end{split} ]
从而就可以得到 $ f‘(x) $的表达式：
[ f'(x) = sum_{i = -P} ^ P b_i sum_{-P leq j leq P, j ot= i} frac{j}{P} prod_{-P leq k leq P, k ot= i, k ot= j} (1 + frac{k}{P} x) ]
当 (x = 1) 时，(i, j) 至少有一个为 (-1) (prod) 里面才不为零，所以也可以暴力计算。

#pragma GCC optimize("2,Ofast,inline")
#include<bits/stdc++.h>
#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define LL long long
#define pii pair<int, int>
#define a(n) a[(n) + 50000]
#define f(n) f[(n) + 50000]
#define g(n) g[(n) + 50000]
using namespace std;
const int N = 1e6 + 10;
const int mod = 998244353;
 
template <typename T> T read(T &x) {
    int f = 0;
    register char c = getchar();
    while (c > '9' || c < '0') f |= (c == '-'), c = getchar();
    for (x = 0; c >= '0' && c <= '9'; c = getchar())
        x = (x << 3) + (x << 1) + (c ^ 48);
    if (f) x = -x;
    return x;
}

namespace Comb {
    const int Maxn = 1e6 + 10;
    
    int fac[Maxn], fav[Maxn], inv[Maxn];
    
    void comb_init() {
        fac[0] = fav[0] = 1;
        inv[1] = fac[1] = fav[1] = 1;
        for (int i = 2; i < Maxn; ++i) {
            fac[i] = 1LL * fac[i - 1] * i % mod;
            inv[i] = 1LL * -mod / i * inv[mod % i] % mod + mod;
            fav[i] = 1LL * fav[i - 1] * inv[i] % mod;
        }
    }
 
    inline int C(int x, int y) {
        if (x < y || y < 0) return 0;
        return 1LL * fac[x] * fav[y] % mod * fav[x - y] % mod;
    }

    inline int Qpow(int x, int p) {
        int ans = 1;
        for (; p; p >>= 1) {
            if (p & 1) ans = 1LL * ans * x % mod;
            x = 1LL * x * x % mod;
        }
        return ans;
    }

    inline int Inv(int x) {
        return Qpow(x, mod - 2);
    }
 
    inline void upd(int &x, int y) {
        (x += y) >= mod ? x -= mod : 0;
    }

    inline int add(int x, int y) {
        if (y < 0) y += mod;
        return (x += y) >= mod ? x - mod : x;
    }

    inline int dec(int x, int y) {
        return (x -= y) < 0 ? x + mod : x;
    }

}
using namespace Comb;

int n, P;
int p[N], f[N], g[N];
int a[N], b[N], s[N];

void Poly_add(int *a, int *f, int p, int k) {
    for (int i = -P; i <= P; ++i) {
        int t = i + p;
        if (t >= -P && t <= P) {
            upd(a(t), k == 1 ? f(i) : mod - f(i));
        }
    }
}

pii calc(int *f) {
    int ans1 = f(-P);
    for (int i = -P + 1; i <= P; ++i) {
        ans1 = 1LL * ans1 * add(1, 1LL * i * inv[P] % mod) % mod;
    }
    int s = 1;
    for (int i = -P + 1; i <= P; ++i) {
        s = 1LL * s * add(1, 1LL * i * inv[P] % mod) % mod;
    }
    int ans2 = 0;
    for (int i = -P; i <= P; ++i) {
        for (int j = -P; j <= ((i == -P) ? P : -P); ++j) {
            if (i == j) continue;
            int t = 1LL * f(i) * j % mod * inv[P] % mod;
            t = 1LL * t * s % mod;
            if (i != -P) {
                t = 1LL * t * Inv(add(1, 1LL * i * inv[P] % mod)) % mod;
            }
            if (j != -P) {
                t = 1LL * t * Inv(add(1, 1LL * j * inv[P] % mod)) % mod;
            }
            if (t < 0) t += mod;
            upd(ans2, t);
        }
    }
    return mp(ans1, ans2);
}

int main() {
    comb_init();
    read(n);
    for (int i = 1; i <= n; ++i) read(s[i]);
    for (int i = 1; i <= n; ++i) read(p[i]), P += p[i];
    f(0) = g(0) = 1;
    for (int i = 1; i <= n; ++i) {
        memset(a, 0, sizeof a);
        Poly_add(a, f, p[i], 1);
        Poly_add(a, f, -p[i], s[i] == 1 ? -1 : 1);
        memcpy(f, a, sizeof f);
        memset(a, 0, sizeof a);
        Poly_add(a, g, p[i], 1);
        Poly_add(a, g, -p[i], 1);
        memcpy(g, a, sizeof g);
    }
    for (int i = 1; i <= P; ++i) {
        swap(f(i), f(-i));
        swap(g(i), g(-i));
    }
    pii F = calc(f);
    pii G = calc(g);
    int ans = dec(1LL * F.se * G.fi % mod, 1LL * F.fi * G.se % mod);
    ans = 1LL * ans * Inv(G.fi) % mod * Inv(G.fi) % mod;
    cout << ans << endl;
    return 0;
}