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

题意

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

思路

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

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

$$\begin{array}{}\text{(1)}& [x^n]f(x)& =\sum _{i_0+i_1+\cdots +i_{w-1}=n-\sum c_i}(\genfrac{}{}{0ex}{}{n}{i_0+c_0,i_1+c_1,\cdots ,i_{w-1}+c_{w-1}})\text{(2)}& & =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{array}$$

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

$$\begin{array}{}\text{(3)}& [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)!}\text{(4)}& \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{array}$$

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

$$\begin{array}{}\text{(5)}& [x^n]F(x)& =[x^n]\prod _{t=0}^{w-1}{g}_t(x)\text{(6)}& & =[x^n]\prod _{t=0}^{w-1}\sum _{j_t=c_t}^n\frac{x^{j_t}}{j_t!}\end{array}$$

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

$$\begin{array}{}\text{(7)}& F(x)& =\prod _{t=0}^{w-1}\sum _{j_t=c_t}^n\frac{x^{j_t}}{j_t!}\end{array}$$

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

$$\begin{array}{}\text{(8)}& F(x)& =\prod _{t=0}^{w-1}(e^x-\sum _{k_t=0}^{c_t-1}\frac{x^{k_t}}{k_t!})\text{(9)}& & =\prod _{t=0}^{w-1}(e^x-G_t(x))\end{array}$$

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

$$\begin{array}{}\text{(10)}& F(x)& =e^{wx}+H_{w-1}\cdot e^{(w-1)x}+\cdots +H_1\cdot e+H_0\end{array}$$

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

$$\begin{array}{}\text{(11)}& An{s}_i& =[x^n]H_i(x)\cdot e^{ix}\text{(12)}& & =\sum _{j=0}^{sizeof(H_i)}[x^j]H_i(x)\cdot [x^{n-j}]e^{ix}\text{(13)}& & =\sum _{j=0}^{sizeof(H_i)}[x^j]H_i(x)\cdot \frac{j^{n-i}}{(n-i)!}\end{array}$$

其中, $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';
    }
}