S2OJ - 1486. 数据结构

题面

题目描述

维护一个正整数多重集合 $S$，初始为空，支持两个操作：

1. 插入：插入一个新数 $x$；
2. 修改：令集合中所有数加 $1$。

每次操作结束后，输出 $S$ 中所有数的 $k$ 次方和，$k$ 预先给定，对 $10^9 + 7$ 取模。

输入格式

第一行两个数 $m, k$，其中 $m$ 表示操作次数。

接下来 $m$ 行，每行可能为以下两种之一：

1. 0 x ，表示插入一个大小为 $x$ 的新元素。
2. 1，表示令集合 $S$ 里所有数加一。

输出格式

输出 $m$ 行，第 $i$ 行表示第 $i$ 次操作结束之后，$S$ 中所有数的 $k$ 次方和。

样例输入输出

样例输入 #1

3 2
0 1
0 1
1

样例输出 #1

1
2
8

数据范围与约定

对于 $100\%$ 的数据，$m \leq 2 \times 10^5$，$1 \leq k \leq 50$，$0 \leq x \leq 10^5$。

思路

这道题考场上大家都是直接把二项式定理拿出来用 A 掉的，复杂度是 $O(nk^2)$，可惜就我一个菜鸡把二项式定理忘掉了，所以通过杨辉三角手推完二项式定理之后顺带着出来了一种 $O(mk)$ 的做法。可惜的是考场上忘记处理负数取模的问题导致猛挂到 20 分。下面带来我的考场思路，可能会有些许罗嗦：

首先题目要求在每次操作完成后输出多重集合 $S$ 中所有数的 $k$ 次方和，因此可以从这里下手。

插入操作的数的 $k$ 次方可以直接加进全局和中，因此只需要考虑如何处理全局加 $1$ 的问题即可。

先写出一些加 $1$ 后 $k$ 次方和的式子进行观察：

1. 当 $k = 1$ 时，

   $$\begin{aligned} & (x_1 + 1)^1 + (x_2 + 1) + \cdots + (x_n + 1) \\ =& 1 \times (x_1 + x_2 + \cdots + x_n) + 1 \times \underbrace{(1 + 1 + \cdots + 1)}_{n \text{ 个 } 1} \\ \end{aligned}$$

2. 当 $k = 2$ 时，

   $$\begin{aligned} & (x_1 + 1)^2 + (x_2 + 1)^2 + \cdots + (x_n + 1)^2 \\ =& (x_1^2 + 2 x_1 + 1) + (x_2^2 + 2 x_2 + 1) + \cdots + (x_n^2 + 2 x_n + 1) \\ =& 1 \times (x_1^2 + x_2^2 + \cdots + x_n^2) + 2 \times (x_1 + x_2 + \cdots + x_n) + 1 \times \underbrace{(1 + 1 + \cdots + 1)}_{n \text{ 个 } 1} \\ \end{aligned}$$

3. 当 $k = 3$ 时，

   $$\begin{aligned} & (x_1 + 1)^3 + (x_2 + 1)^3 + \cdots + (x_n + 1)^3 \\ =& (x_1^3 + 3 x_1^2 + 3 x_1 + 1) + (x_2^3 + 3 x_2^2 + 3 x_2 + 1) + \cdots + (x_n^3 + 3 x_n^2 + 3 x_n + 1) \\ =& 1 \times (x_1^3 + x_2^3 + \cdots + x_n^3) + 3 \times (x_1^2 + x_2^2 + \cdots + x_n^2) + 3 \times (x_1 + x_2 + \cdots + x_n) + 1 \times \underbrace{(1 + 1 + \cdots + 1)}_{n \text{ 个 } 1} \\ \end{aligned}$$

4. 当 $k = 4$ 时，

   $$\begin{aligned} & (x_1 + 1)^4 + (x_2 + 1)^4 + \cdots + (x_n + 1)^4 \\ =& (x_1^4 + 4 x_1^3 + 6 x_1^2 + 4 x_1 + 1) + (x_2^4 + 4 x_2^3 + 6 x_2^2 + 4 x_2 + 1) + \cdots + (x_n^4 + 4 x_n^3 + 6 x_n^2 + 4 x_n + 1) \\ =& 1 \times (x_1^4 + x_2^4 + \cdots + x_n^4) + 4 \times (x_1^3 + x_2^3 + \cdots + x_n^3) + 6 \times (x_1^2 + x_2^2 + \cdots + x_n^2) + 4 \times (x_1 + x_2 + \cdots + x_n) + 1 \times \underbrace{(1 + 1 + \cdots + 1)}_{n \text{ 个 } 1} \\ \end{aligned}$$

可以观察到系数呈现出了杨辉三角的形式：

$$\begin{array}{c} 1 \\ 1 \quad 1 \\ 1 \quad 2 \quad 1 \\ 1 \quad 3 \quad 3 \quad 1 \\ 1 \quad 4 \quad 6 \quad 4 \quad 1 \\ \end{array}$$

那么可以考虑记录全局累计增加的数 $\Delta$，然后每次全局加之后都将 $\Delta + 1$，再重新计算全局和。此时需要将新插入的数 $x$ 作为 $x - \Delta$ 存入以便计算。

现在再来看一看 $(x_i + \Delta)^k$ 的问题，有一种想法是可以将其按照 $((x_i + (\Delta - 1)) + 1)^k$ 的形式递归展开，比如 $(x + 2)^2$ 就可以这样展开：

$$\begin{aligned} & (x + 2)^2 \\ =& (x + 1)^2 + 2 (x + 1) + 1 \\ =& (x^2 + 2x + 1) + 2 (x + 1) + 1 \\ =& x^2 + 4x + 4 \\ \end{aligned}$$

当 $k = 2$ 时，系数分布如下图所示：

$$\begin{array}{c} 1 \times 2^0 \\ 1 \times 2^0 \quad 1 \times 2^1 \\ 1 \times 2^0 \quad 2 \times 2^1 \quad 1 \times 2^2 \\ 1 \times 2^0 \quad 3 \times 2^1 \quad 3 \times 2^2 \quad 1 \times 2^3 \\ 1 \times 2^0 \quad 4 \times 2^1 \quad 6 \times 2^2 \quad 4 \times 2^3 \quad 1 \times 2^4 \\ \end{array}$$

推广得：

$$\begin{array}{c} 1 \times k^0 \\ 1 \times k^0 \quad 1 \times k^1 \\ 1 \times k^0 \quad 2 \times k^1 \quad 1 \times k^2 \\ 1 \times k^0 \quad 3 \times k^1 \quad 3 \times k^2 \quad 1 \times k^3 \\ 1 \times k^0 \quad 4 \times k^1 \quad 6 \times k^2 \quad 4 \times k^3 \quad 1 \times k^4 \\ \end{array}$$

举个稍大一些的例子：

$$\begin{aligned} & (x_1 + 2)^4 + (x_2 + 2)^4 \\ =& (x_1^4 + 8 x_1^3 + 24 x_1^2 + 32 x_1 + 2^4) + (x_2^4 + 8 x_2^3 + 24 x_2^2 + 32 x_2 + 2^4) \\ =& (x_1^4 + x_2^4) + 8 \times (x_1^3 + x_2^3) + 24 \times (x_1^2 + x_2^2) + 32 \times (x_1 + x_2) + (2^4 + 2^4) \\ \end{aligned}$$

设左对齐的杨辉三角的第 $i$ 行第 $j$ 列上的数为 $f_{i, j}$，那么可以将总和记为：

$$\sum_{i = 0}^{k} (f_{k, i} \times \Delta^i \times \sum_{j = 1}^{n} x_j^{k - i})$$

显然 $\sum_{j = 1}^{n} x_j^{k - i}$ 可以使用前缀和优化到 $O(1)$，记为 $\mathit{sum}$，最终形式为：

$$\sum_{i = 0}^{k} (f_{k, i} \times \Delta^i \times \mathit{sum}^{k - i})$$

复杂度为 $O(mk)$。

代码

#include <iostream>

using std::cin;
using std::cout;
const char endl = '\n';

const int N = 2e5 + 5,
          K = 55;
const int mod = 1e9 + 7;

int m, k, d, cnt;
long long f[K][K], s[K], sum;

long long binpow(long long a, long long b) {
    a %= mod;

    long long res = 1;

    while (b) {
        if (b & 1) res = res * a % mod;
        a = a * a % mod;
        b >>= 1;
    }

    return res;
}

int main() {
    std::ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> m >> k;

    f[0][0] = 1;
    for (int i = 1; i <= k; i++) {
        f[i][0] = 1;

        for (int j = 1; j <= i; j++) {
            f[i][j] = (f[i - 1][j - 1] + f[i - 1][j]) % mod;
        }
    }

    while (m--) {
        int op;

        cin >> op;

        if (op == 0) {
            int x;
            long long p = 1;

            cin >> x;

            s[0] = (s[0] + 1) % mod;

            for (int i = 1; i <= k; i++) {
                p = (p * (x - d) % mod + mod) % mod;
                s[i] = (s[i] + p) % mod;
            }

            sum = (sum + binpow(x, k)) % mod;
        } else {  // op == 1
            d++;

            sum = 0;
            for (int i = 0; i <= k; i++) {
                sum = (sum + f[k][i] * binpow(d, i) % mod * s[k - i] % mod) % mod;
            }
        }

        cout << sum << endl;
    }

    return 0;
}