[ZJOI2019] 开关

\(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) \]
证明可以暴力展开:
\[ \begin{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 \not= i} \frac{j}{P} \prod_{-P \leq k \leq P, k \not= i, k \not= 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;
}

猜你喜欢

转载自www.cnblogs.com/Vexoben/p/11756605.html