2022 牛客多校 第四场 C - Easy Counting Problem

题意

给定 \(0\) 到 \(w-1\) ( \(w\leq 10\) ) 每个数至少出现的次数 ( 合起来不超过 \(50000\) ), \(q\) ( \(q\leq 300\) ) 次询问, 每次要求计算长度为 \(n\) ( \(n\leq 10^7\) ) 的只包含 \(0\) 到 \(w-1\) 的数串且满足以上条件的方案数

思路

显然, 对于每个数的出现次数, 以串长作为生成函数的指数, 有

( \([x^n]f(x)\) 表示生成函数的第 \(n\) 项系数 )

\[\begin{align} [x^n] f(x) &= \sum_{i_0+i_1+\cdots +i_{w-1}=n-\sum c_i} \binom{n}{i_0+c_0, i_1+c_1,\cdots,i_{w-1}+c_{w-1}} \\ &= n! \sum_{i_0+i_1+\cdots +i_{w-1}=n-\sum c_i} \prod_{t=0}^{w-1} \frac{1}{(i_t+c_t)!} \\ \end{align} \]

考虑到这里的 \(\sum\) 其实就是 \(w\) 个 \(for\) 循环, 对它进行展化简得到

\[\begin{align} [x^n] f(x) &= n! \prod_{t=0}^{w-1} [\sum i_t + c_t == n] \sum_{i_t=0}^{n-c_t} \frac{1}{(i_t + c_t)!} \\ \frac{[x^n] f(x)}{n!} &= \prod_{t=0}^{w-1} [\sum j_t == n] \sum_{j_t=c_t}^{n} \frac{1}{j_t!} \end{align} \]

上面式子的 \([\sum j_t == n]\) 表示累和等于 \(n\) 情况下才有贡献, 对每个 \(i_t\) 进行展开, 记 \(F(x)\) 表示 \([x^n]F(x)\), 并将右边也写作 \([x^n]g(x)\) 的形式

\[\begin{align} [x^n]F(x) &= [x^n] \prod_{t=0}^{w-1} g_t(x) \\ &= [x^n] \prod_{t=0}^{w-1} \sum_{j_t=c_t}^{n} \frac{x^{j_t}}{j_t!} \end{align} \]

上式表示, 对每个数 \(t\) 取 \(j_t\) 个, 最后生成函数乘起来的 \([x^n]\) 项系数即为结果, 那么接下来可以去掉 \([x^n]\) 了

\[\begin{align} F(x) &= \prod_{t=0}^{w-1} \sum_{j_t=c_t}^{n} \frac{x^{j_t}}{j_t!} \end{align} \]

显然枚举到 \(n\) 实在是太大了, 此时注意到有点像函数 \(e^x\) 的泰勒展开形式, 那么就有

\[\begin{align} F(x) &= \prod_{t=0}^{w-1} (e^x - \sum_{k_t=0}^{c_t-1} \frac{x^{k_t}}{k_t!}) \\ &= \prod_{t=0}^{w-1} (e^x - G_t(x)) \end{align} \]

其中, 记 \(G_t(x) = \sum_{k_t=0}^{c_t-1} \frac{x^{k_t}}{k_t!}\), 那么接下来进行展开, 简单记各个生成函数 \(G_t\) 乘积的结果为 \(H_m\), 故 \(F(x)\) 形式如下

\[\begin{align} F(x) &= e^{wx} + H_{w-1}\cdot e^{(w-1)x} + \cdots + H _1\cdot e + H_0 \end{align} \]

计算各个 \(H_m(x)\), 只需要对 \(G_t(x)\) 状压乘起来就行, 接下来对于每个项, 有

\[\begin{align} Ans_i &= [x^n] H_i(x) \cdot e^{ix} \\ &= \sum_{j=0} ^{sizeof(H_i)} [x^j]H_i(x) \cdot [x^{n-j}]e^{ix} \\ &= \sum_{j=0} ^{sizeof(H_i)} [x^j]H_i(x) \cdot \frac{j^{n-i}}{(n-i)!} \\ \end{align} \]

其中, \(O(n)\) 预处理阶乘, 消费空间 \(40MB\) , 每个多项式长度不会超过 \(50000\) , 询问 \(300\) 次, 一共最多有 \(10\) 个多项式, 复杂度在 \(1.5 \cdot 10^8\), 非常极限, 如果继续使用快速幂, 很显然会超时, 可以预处理出 \(w\) 个低 \(14\) 位和高 \(14\) 位的幂次乘积, 从而 \(O(1)\) 得到幂次值

代码

代码中本来想滚动优化幂次值, 结果发现发现直接预处理 \(14\) 位幂次就行了, 原来的离线处理询问没有删干净, 但是懒得改了, 所以有离线处理

constexpr int N = 10000010;
constexpr int sq = 4096;

void sol() {
    int w;
    std::cin >> w;
    std::vector<int> c(w);
    for(int i=0;i<w;++i) {
        std::cin >> c[i];
    }

    std::vector<int> fac(N + 1, 1), ifac(N + 1, 1);

    for(int i=2;i<=N;++i) {
        fac[i] = mul(fac[i - 1], i);
    }
    ifac[N] = Poly::power(fac[N], MOD - 2);
    for(int i=N-1;i>=2;--i) {
        ifac[i] = mul(ifac[i + 1], i + 1);
    }

    std::vector<poly> a(1 << w), d(w + 1);
    for(int T = 0; T < w; ++T) {
        int t = 1 << T;
        a[t].resize(c[T]);
        for(int i = 0; i < c[T]; ++i) {
            a[t][i] = add(0, MOD - ifac[i]);
        }
    }

    a[0] = poly(1, 1);
    for(int dp=0;dp<1<<w;++dp) {
        for(int j=0;j<w;++j)
        if((dp & 1 << j) == 0 && a[dp | 1 << j].empty()) {
            a[dp | 1 << j] = Poly::mul(a[dp], a[1 << j]);
        }
    }
    for(int dp=0;dp<1<<w;++dp) {
        int cto = 0;
        for(int tmp = dp; tmp; tmp >>= 1) if(tmp&1) ++cto;
        int ctz = w - cto;
        d[ctz].resize(std::max(a[dp].size(), d[ctz].size()));
        for(int i=0;i<a[dp].size();++i) d[ctz][i] = add(d[ctz][i], a[dp][i]);
        poly().swap(a[dp]);
    }

    int Q;
    std::cin >> Q;
    std::vector<int> q(Q),pos(Q);
    for(int i=0;i<Q;++i) {
        std::cin >> q[i];
    }
    std::iota(pos.begin(), pos.end(), 0);
    std::sort(pos.begin(), pos.end(), [&](int x,int y) {return q[x] > q[y];});

    std::vector<std::vector<std::array<int,2>>> tb(w + 1, std::vector<std::array<int,2>>(sq));
    for(int i=0;i<sq;++i) tb[1][i][0] = tb[1][i][1] = 1;
    for(int t=2;t<=w;++t) {
        int qwq = Poly::power(t, sq);
        tb[t][0][0] = tb[t][0][1] = 1;
        for(int i=1;i<sq;++i) {
            tb[t][i][0] = mul(tb[t][i - 1][0], t);
            tb[t][i][1] = mul(tb[t][i - 1][1], qwq);
        }
    }

    for(int I=0;I<Q;++I) {
        int n = q[pos[I]];
        int res = 0;
 
        if(n < (int)d[0].size()) res = add(res, d[0][n]);
        for(int t=1;t<=w;++t) {
            int m = (int)d[t].size();
            for(int i=0;i<m&&i<=n;++i) {
                res = add(res, mul(mul(ifac[n - i], mul(tb[t][(n - i) % sq][0], tb[t][(n - i) / sq][1])), d[t][i]));
            }
        }

        q[pos[I]] = mul(res, fac[n]);
    }
    for(int e : q) {
        std::cout << e << '\n';
    }
}