[2022 牛客多校4 C] Easy Counting Problem (生成函数 NTT)

题意

给定一个正整数 w w w w w w 个数 c 0 , c 1 , ⋯   , c w − 1 c_0, c_1, \cdots,c_{w - 1} c0,c1,,cw1

q q q 组询问,每次询问给定一个正整数 n n n,计算有多少个长度为 n n n 的字符串满足:

  • 每个字符只能取数字 0 ∼ w − 1 0 \sim w - 1 0w1
  • 数字 i i i 至少出现 c i c_i ci

998   244   353 998 \, 244 \, 353 998244353 取模。

2 ≤ w ≤ 10 , 1 ≤ c i ≤ 5 × 1 0 4 , ∑ i = 0 w − 1 c i ≤ 5 × 1 0 4 2 \le w \le 10, 1 \le c_i \le 5 \times 10 ^ 4, \sum\limits_{i = 0} ^ {w - 1}c_i \le 5 \times 10 ^ 4 2w10,1ci5×104,i=0w1ci5×104

1 ≤ q ≤ 300 , 1 ≤ n ≤ 1 0 7 1 \le q \le 300, 1 \le n \le 10 ^ 7 1q300,1n107

分析:

首先我们可以写出每个数字 i i i EGF \textbf{EGF} EGF
∑ j = c i ∞ x j j ! \sum_{j = c_i} ^ {\infty} \frac{x ^ j}{j!} j=cij!xj
那么每个数字的 EGF \textbf{EGF} EGF 做乘积表示满足条件的所有长度的字符串的方案数
∏ i = 0 w − 1 ∑ j = c i ∞ x j j ! \prod_{i = 0} ^ {w - 1}\sum_{j = c_i} ^ {\infty} \frac{x ^ j}{j!} i=0w1j=cij!xj
可以把和式用前缀和相减拆一下 ∑ j = c i ∞ x j j ! = ∑ j = 0 ∞ x j j ! − ∑ j = 0 c i − 1 x j j ! \sum\limits_{j = c_i} ^ {\infty} \dfrac{x ^ j}{j!} = \sum\limits_{j = 0} ^ {\infty} \dfrac{x ^ j}{j!} - \sum\limits_{j = 0} ^ {c_i - 1} \dfrac{x ^ j}{j!} j=cij!xj=j=0j!xjj=0ci1j!xj,发现第一项为 e x e ^ x ex,故答案为
∏ i = 0 w − 1 ( e x − ∑ j = 0 c i − 1 x j j ! ) \prod_{i = 0} ^ {w - 1}(e ^ x - \sum_{j = 0} ^ {c_i - 1} \frac{x ^ j}{j!}) i=0w1(exj=0ci1j!xj)
由于 w ≤ 10 w \le 10 w10,所以考虑暴力展开式子,做换元 e x → y e ^ x \rightarrow y exy

在展开式子的过程中,假设当前的多项式为 f = A 0 + A 1 y + A 2 y 2 + A 3 y 3 + ⋯ f = A_0 + A_1 y + A_2y ^ 2 + A_3y^3 + \cdots f=A0+A1y+A2y2+A3y3+,那么新遇到一个多项式 ( y + g i ) (y + g_i) (y+gi) 其中 g i = − ∑ j = 0 c i − 1 x j j ! g_i = -\sum\limits_{j = 0} ^ {c_i - 1} \dfrac{x ^ j}{j!} gi=j=0ci1j!xj, 则结果变为 f ∗ y + f ∗ g i f * y + f * g_i fy+fgi ( ∗ * 表示多项式卷积),前一项为 A 0 y + A 1 y 2 + A 2 y 3 + A 3 y 4 + ⋯ A_0y + A_1 y ^ 2 + A_2y ^ 3 + A_3y^4 + \cdots A0y+A1y2+A2y3+A3y4+,那么后一项是 f f f 的每一项系数与 g i g_i gi 的多项式卷积,为 A 0 ∗ g i + ( A 1 ∗ g i ) y + ( A 2 ∗ g i ) y 2 + ( A 3 ∗ g i ) y 3 + ⋯ A_0 * g_i + (A_1 * g_i) y + (A_2 * g_i)y ^ 2 + (A_3 * g_i)y^3 + \cdots A0gi+(A1gi)y+(A2gi)y2+(A3gi)y3+,那么答案就为
A 0 ∗ g i + ( A 0 + A 1 ∗ g i ) y + ( A 1 + A 2 ∗ g i ) y 2 + ( A 2 + A 3 ∗ g i ) y 3 + ⋯ A_0 * g_i + (A_0 + A_1 * g_i)y + (A_1 + A_2*g_i)y ^ 2 + (A_2 + A_3*g_i)y ^ 3 + \cdots A0gi+(A0+A1gi)y+(A1+A2gi)y2+(A2+A3gi)y3+
这样就预处理好了总答案,现考虑回答每组询问,我们知道最后的答案是形如 ∑ i = 0 w − 1 e i x F i ( x ) \sum\limits_{i = 0} ^ {w - 1} e ^ {ix} F_i(x) i=0w1eixFi(x) 的多项式,我们需要知道每一项的第 n n n 项系数,由于 ∑ i = 0 w − 1 c i ≤ 5 × 1 0 4 \sum\limits_{i = 0} ^ {w - 1}c_i \le 5 \times 10 ^ 4 i=0w1ci5×104,我们可以在询问里对于每个 i i i 直接枚举 F i ( x ) F_i(x) Fi(x) 的项数,设当前枚举到了第 j j j 项,那么需要在 e i x e ^ {ix} eix 中取出第 n − j n - j nj 项,也就是 e i x = 1 + ( i x ) 1 1 ! + ( i x ) 2 2 ! + ( i x ) 3 3 ! + ⋯ e ^ {ix} = 1 + \dfrac{(ix) ^ 1}{1!} + \dfrac{(ix) ^ 2}{2!} + \dfrac{(ix) ^ 3}{3!} + \cdots eix=1+1!(ix)1+2!(ix)2+3!(ix)3+ 的第 n − j n - j nj 项,为 i n − j ( n − j ) ! \dfrac{i ^ {n - j}}{(n - j)!} (nj)!inj

那么答案就为
n ! × ∑ i = 0 w ∑ j = 0 min ⁡ ( n , ∣ F i ( x ) ∣ ) [ x j ] F i ( x ) × i n − j ( n − j ) ! n! \times \sum_{i = 0} ^ {w} \sum_{j = 0} ^ {\min(n, |F_i(x)|)} [x ^ j] F_i(x) \times \frac{i ^ {n - j}}{(n - j)!} n!×i=0wj=0min(n,Fi(x))[xj]Fi(x)×(nj)!inj

代码:

#include <bits/stdc++.h>
using namespace std;
using i64 = long long;
constexpr int mod = 998244353;
int norm(int x) {
    if (x < 0) {
        x += mod;
    }
    if (x >= mod) {
        x -= mod;
    }
    return x;
}
template<class T>
T power(T a, int b) {
    T res = 1;
    for (; b; b /= 2, a *= a) {
        if (b % 2) {
            res *= a;
        }
    }
    return res;
}
struct Z {
    int x;
    Z(int x = 0) : x(norm(x)) {}
    int val() const {
        return x;
    }
    Z operator-() const {
        return Z(norm(mod - x));
    }
    Z inv() const {
        assert(x != 0);
        return power(*this, mod - 2);
    }
    Z &operator*=(const Z &rhs) {
        x = i64(x) * rhs.x % mod;
        return *this;
    }
    Z &operator+=(const Z &rhs) {
        x = norm(x + rhs.x);
        return *this;
    }
    Z &operator-=(const Z &rhs) {
        x = norm(x - rhs.x);
        return *this;
    }
    Z &operator/=(const Z &rhs) {
        return *this *= rhs.inv();
    }
    friend Z operator*(const Z &lhs, const Z &rhs) {
        Z res = lhs;
        res *= rhs;
        return res;
    }
    friend Z operator+(const Z &lhs, const Z &rhs) {
        Z res = lhs;
        res += rhs;
        return res;
    }
    friend Z operator-(const Z &lhs, const Z &rhs) {
        Z res = lhs;
        res -= rhs;
        return res;
    }
    friend Z operator/(const Z &lhs, const Z &rhs) {
        Z res = lhs;
        res /= rhs;
        return res;
    }
    friend istream &operator>>(istream &is, Z &a) {
        i64 v;
        is >> v;
        a = Z(v);
        return is;
    }
    friend ostream &operator<<(ostream &os, const Z &a) {
        return os << a.val();
    }
};
vector<int> rev;
vector<Z> roots{0, 1};
void dft(vector<Z> &a) {
    int n = a.size();
    if (int(rev.size()) != n) {
        int k = __builtin_ctz(n) - 1;
        rev.resize(n);
        for (int i = 0; i < n; i ++) {
            rev[i] = rev[i >> 1] >> 1 | (i & 1) << k;
        }
    }
    for (int i = 0; i < n; i ++) {
        if (rev[i] < i) {
            swap(a[i], a[rev[i]]);
        }
    }
    if (int(roots.size()) < n) {
        int k = __builtin_ctz(roots.size());
        roots.resize(n);
        while ((1 << k) < n) {
            Z e = power(Z(3), (mod - 1) >> (k + 1));
            for (int i = 1 << (k - 1); i < (1 << k); i ++) {
                roots[i << 1] = roots[i];
                roots[i << 1 | 1] = roots[i] * e;
            }
            k ++;
        }
    }
    for (int k = 1; k < n; k *= 2) {
        for (int i = 0; i < n; i += 2 * k) {
            for (int j = 0; j < k; j ++) {
                Z u = a[i + j], v = a[i + j + k] * roots[k + j];
                a[i + j] = u + v, a[i + j + k] = u - v;
            }
        }
    }
}
void idft(vector<Z> &a) {
    int n = a.size();
    reverse(a.begin() + 1, a.end());
    dft(a);
    Z inv = (1 - mod) / n;
    for (int i = 0; i < n; i ++) {
        a[i] *= inv;
    }
}
struct Poly {
    vector<Z> a;
    Poly() {}
    Poly(const vector<Z> &a) : a(a) {}
    Poly(const initializer_list<Z> &a) : a(a) {}
    int size() const {
        return a.size();
    }
    void resize(int n) {
        a.resize(n);
    }
    Z operator[](int idx) const {
        if (idx < size()) {
            return a[idx];
        } else {
            return 0;
        }
    }
    Z &operator[](int idx) {
        return a[idx];
    }
    Poly mulxk(int k) const {
        auto b = a;
        b.insert(b.begin(), k, 0);
        return Poly(b);
    }
    Poly modxk(int k) const {
        k = min(k, size());
        return Poly(vector<Z>(a.begin(), a.begin() + k));
    }
    Poly divxk(int k) const {
        if (size() <= k) {
            return Poly();
        }
        return Poly(vector<Z>(a.begin() + k, a.end()));
    }
    friend Poly operator+(const Poly &a, const Poly &b) {
        vector<Z> res(max(a.size(), b.size()));
        for (int i = 0; i < int(res.size()); i ++) {
            res[i] = a[i] + b[i];
        }
        return Poly(res);
    }
    friend Poly operator-(const Poly &a, const Poly &b) {
        vector<Z> res(max(a.size(), b.size()));
        for (int i = 0; i < int(res.size()); i ++) {
            res[i] = a[i] - b[i];
        }
        return Poly(res);
    }
    friend Poly operator*(Poly a, Poly b) {
        if (a.size() == 0 || b.size() == 0) {
            return Poly();
        }
        int sz = 1, tot = a.size() + b.size() - 1;
        while (sz < tot) {
            sz *= 2;
        }
        a.a.resize(sz);
        b.a.resize(sz);
        dft(a.a);
        dft(b.a);
        for (int i = 0; i < sz; i ++) {
            a.a[i] = a[i] * b[i];
        }
        idft(a.a);
        a.resize(tot);
        return a;
    }
    friend Poly operator*(Z a, Poly b) {
        for (int i = 0; i < int(b.size()); i ++) {
            b[i] *= a;
        }
        return b;
    }
    friend Poly operator*(Poly a, Z b) {
        for (int i = 0; i < int(a.size()); i ++) {
            a[i] *= b;
        }
        return a;
    }
    Poly &operator+=(Poly b) {
        return (*this) = (*this) + b;
    }
    Poly &operator-=(Poly b) {
        return (*this) = (*this) - b;
    }
    Poly &operator*=(Poly b) {
        return (*this) = (*this) * b;
    }
    Poly deriv() const {
        if (a.empty()) {
            return Poly();
        }
        vector<Z> res(size() - 1);
        for (int i = 0; i < size() - 1; i ++) {
            res[i] = (i + 1) * a[i + 1];
        }
        return Poly(res);
    }
    Poly integr() const {
        vector<Z> res(size() + 1);
        for (int i = 0; i < size(); i ++) {
            res[i + 1] = a[i] / (i + 1);
        }
        return Poly(res);
    }
    Poly inv(int m) const {
        Poly x{a[0].inv()};
        int k = 1;
        while (k < m) {
            k *= 2;
            x = (x * (Poly{2} - modxk(k) * x)).modxk(k);
        }
        return x.modxk(m);
    }
    Poly log(int m) const {
        return (deriv() * inv(m)).integr().modxk(m);
    }
    Poly exp(int m) const {
        Poly x{1};
        int k = 1;
        while (k < m) {
            k *= 2;
            x = (x * (Poly{1} - x.log(k) + modxk(k))).modxk(k);
        }
        return x.modxk(m);
    }
    Poly pow(int k, int m) const {
        int i = 0;
        while (i < size() && a[i].val() == 0) {
            i ++;
        }
        if (i == size() || 1LL * i * k >= m) {
            return Poly(vector<Z>(m));
        }
        Z v = a[i];
        auto f = divxk(i) * v.inv();
        return (f.log(m - i * k) * k).exp(m - i * k).mulxk(i * k) * power(v, k);
    }
    Poly sqrt(int m) const {
        Poly x{1};
        int k = 1;
        while (k < m) {
            k *= 2;
            x = (x + (modxk(k) * x.inv(k)).modxk(k)) * ((mod + 1) / 2);
        }
        return x.modxk(m);
    }
    Poly mulT(Poly b) const {
        if (b.size() == 0) {
            return Poly();
        }
        int n = b.size();
        reverse(b.a.begin(), b.a.end());
        return ((*this) * b).divxk(n - 1);
    }
};
vector<Z> fact, infact;
void init(int n) {
    fact.resize(n + 1), infact.resize(n + 1);
    fact[0] = infact[0] = 1;
    for (int i = 1; i <= n; i ++) {
        fact[i] = fact[i - 1] * i;
    }
    infact[n] = fact[n].inv();
    for (int i = n; i; i --) {
        infact[i - 1] = infact[i] * i;
    }
}
signed main() {
    init(1e7);
    cin.tie(0) -> sync_with_stdio(0);
    int w;
    cin >> w;
    vector<Poly> ans(w + 1);
    ans[0] = {1};
    for (int i = 1; i <= w; i ++) {
        int c;
        cin >> c;
        vector<Z> g(c);
        for (int j = 0; j < c; j ++) {
            g[j] = -infact[j];
        }
        for (int j = i; j; j --) {
            ans[j] = ans[j] * Poly(g) + ans[j - 1];
        }
        ans[0] = ans[0] * Poly(g);
    }
    int m;
    cin >> m;
    while (m --) {
        int n;
        cin >> n;
        Z res;
        for (int i = 0; i <= w; i ++) {
            int v = min(ans[i].size() - 1, n);
            Z Pow = power(Z(i), n - v);
            for (int j = v; ~j; j --) {
                res += ans[i][j] * Pow * infact[n - j];
                Pow *= i;
            }
        }
        cout << res * fact[n] << "\n";
    }
}
  • 4
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值