E. Thief in a a Shop - 题解

比赛与标签

比赛: Educational Codeforces Round 9 标签: divide and conquer, dp, fft, math 难度: *2400

题目大意喵~

有一位小偷先生要去商店“进货”啦！他有一个神奇的背包，正好能装下 k 件物品。商店里有 n 种商品，每种商品都有无限件，第 i 种商品的价值是 a_i。

这位小偷先生非常贪心（也很讲究！），他不多不少，一定会拿走恰好 k 件商品。他可以拿多件同种类的商品哦。

现在，请你帮他算一算，他拿走的这 k 件商品，所有可能的总价值是多少呢？把这些可能的总价值按从小到大的顺序输出就可以啦，喵~

输入格式:

• 第一行是两个整数 n 和 k。
• 第二行是 n 个整数 a_i，代表每种商品的价值。

输出格式:

• 一行，所有可能的总价值，用空格隔开，并按升序排列。

解题思路，来一起探险吧！

看到这个问题，大家的第一反应是不是经典的动态规划（DP）呢？比如说，dp[i][j] 表示拿 i 件物品能否凑出总价值 j。但是我们来分析一下数据范围：k 和 a_i 都是 1000，那么总价值最高可以达到 k * max(a_i) = 1000 * 1000 = 10^6。DP的状态数量会是 k * (k * max(a_i))，这实在是太大了，会超时超到天上去的，不行不行！(>ω<)

那么，我们得换个思路啦！这种“求若干个数相加能凑出的所有和”的问题，有一个超级厉害的武器——生成函数 (Generating Function)，也叫母函数！

1. 生成函数的魔法 ✨

让我们把每种商品的选择看作一个多项式。对于一种价值为 c 的商品，我们可以用 x^c 这一项来代表它。那么，如果我们有价值为 a_1, a_2, ..., a_n 的商品，我们可以构造一个“选择多项式” P(x)：

P(x) = x^{a_1} + x^{a_2} + ... + x^{a_n}

这个多项式代表了我们拿一件商品的所有可能性。x 的指数就代表了这件商品的价值。

那如果我们拿两件呢？比如我们拿了价值为 a_i 和 a_j 的两件商品，总价值就是 a_i + a_j。在多项式里，这对应着 x^{a_i} * x^{a_j} = x^{a_i + a_j}。

所以，要得到拿两件商品所有可能的总价值，我们只需要计算 P(x) * P(x) = P(x)^2！展开后的多项式中，所有 x 的指数就是所有可能的总价值和。

以此类推，我们要拿 k 件商品，那不就是计算 P(x)^k 吗？最后，我们只要看看 P(x)^k 这个多项式里，哪些项的系数不为零，这些项的指数就是我们要求的答案啦！是不是很神奇，喵？

2. 加速！FFT 冲呀！🚀

直接计算多项式乘法还是很慢的。但是，我们有 快速傅里叶变换 (Fast Fourier Transform, FFT) 这个大杀器！FFT 可以在 O(N log N) 的时间内完成两个 N 次多项式的乘法，其中 N 是结果多项式的最高次数。

所以，我们的计划是：

1. 构造初始多项式 P(x)。
2. 使用快速幂（也叫二进制幂） 的思想来计算 P(x)^k。这需要 O(log k) 次多项式乘法。
3. 每次多项式乘法都用 FFT 来加速。

3. 一个小小的优化技巧 🐾

直接用 a_i 来构造多项式，最高次会达到 k * max(a_i)，大约是 10^6。虽然可行，但我们可以做得更好！

注意到，如果我们把所有商品的价格都减去同一个值，比如最小价格 min_a，那么总价格的相对关系是不变的。 假设我们拿了 k 件商品，它们的原始价格是 c_1, c_2, ..., c_k。 总价 S = c_1 + c_2 + ... + c_k。 如果我们把每个价格都减去 min_a，得到新价格 b_i = c_i - min_a。 那么新的总价 S' = b_1 + b_2 + ... + b_k = (c_1 - min_a) + ... + (c_k - min_a) = S - k * min_a。 所以 S = S' + k * min_a。

这意味着，我们可以先用变换后的价格 b_i = a_i - min_a 来解决问题，求出所有可能的新总价 S'，最后再给每个 S' 加上 k * min_a 就好啦！

这样做的好处是，新的价格 b_i 从 0 开始，最大值也变小了 (max(a_i) - min(a_i))，多项式的最高次也相应减小，计算起来更舒服~

总结一下我们的最终方案：

1. 读入所有价格 a_i，去重并排序。
2. 找到最小价格 min_a。
3. 将所有价格 a_i 转化为 b_i = a_i - min_a。
4. 根据 b_i 构造一个 01 系数的多项式 P(x)，如果 v 是一个有效的 b_i 值，那么 P(x) 中 x^v 的系数为 1，否则为 0。
5. 使用 FFT 和快速幂计算出结果多项式 R(x) = P(x)^k。
6. 遍历 R(x) 的系数。如果 x^i 的系数大于 0，说明 i 是一个可能的“新总价”，那么 i + k * min_a 就是一个最终答案！

好啦，思路清晰了，让我们看看代码是怎么实现的吧！

代码实现，一起来敲代码喵！

#include <iostream>
#include <vector>
#include <algorithm>
#include <complex>
#include <cmath>

// 快速 I/O，让程序跑得像猫一样快！
void fast_io() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(nullptr);
}

// 定义复数类型，FFT 的基础喵
using base = std::complex<double>;
const double PI = acos(-1.0);

// FFT 核心函数，实现多项式从系数表示到点值表示的转换（以及逆转换）
void fft(std::vector<base>& a, bool invert) {
    int n = a.size();
    if (n <= 1) return;

    // 蝴蝶变换前的“位逆序置换”，让数据排好队~
    for (int i = 1, j = 0; i < n; i++) {
        int bit = n >> 1;
        for (; j & bit; bit >>= 1)
            j ^= bit;
        j ^= bit;
        if (i < j)
            std::swap(a[i], a[j]);
    }

    // 从底层向上，一层层进行蝶形运算
    for (int len = 2; len <= n; len <<= 1) {
        double ang = 2 * PI / len * (invert ? -1 : 1); // 根据正向或逆向FFT计算单位根
        base wlen(cos(ang), sin(ang));
        for (int i = 0; i < n; i += len) {
            base w(1);
            for (int j = 0; j < len / 2; j++) {
                base u = a[i + j];
                base v = a[i + j + len / 2] * w;
                a[i + j] = u + v;
                a[i + j + len / 2] = u - v;
                w *= wlen;
            }
        }
    }

    // 如果是逆FFT (IFFT)，需要将结果除以 n
    if (invert) {
        for (base& x : a) {
            x /= n;
        }
    }
}

// 使用 FFT 加速的多项式乘法
std::vector<long long> multiply(const std::vector<long long>& a, const std::vector<long long>& b) {
    if (a.empty() || b.empty()) {
        return {};
    }
    std::vector<base> fa(a.begin(), a.end()), fb(b.begin(), b.end());
    int n = 1;
    // 计算FFT需要的长度，必须是2的幂且大于等于结果多项式的项数
    while (n < a.size() + b.size() - 1) n <<= 1;
    fa.resize(n);
    fb.resize(n);

    // 1. 正向 FFT
    fft(fa, false);
    fft(fb, false);

    // 2. 点值相乘
    for (int i = 0; i < n; i++) {
        fa[i] *= fb[i];
    }
    
    // 3. 逆向 FFT
    fft(fa, true);

    std::vector<long long> result(a.size() + b.size() - 1);
    for (size_t i = 0; i < result.size(); i++) {
        // 由于浮点数精度问题，需要四舍五入取整
        result[i] = round(fa[i].real());
    }
    return result;
}

// 多项式快速幂，计算 p(x)^exp
std::vector<long long> poly_pow(std::vector<long long> p, int exp) {
    std::vector<long long> res = {1}; // 结果初始化为 p(x)^0 = 1
    while (exp > 0) {
        if (exp % 2 == 1) {
            res = multiply(res, p);
            // 优化：我们只关心某个价值是否可能凑出，不关心有多少种方法
            // 所以系数大于0的都设为1，可以防止数值过大和减小精度误差
            for(auto& val : res) if (val > 0) val = 1;
        }
        if (exp > 1) {
            p = multiply(p, p);
            for(auto& val : p) if (val > 0) val = 1;
        }
        exp /= 2;
    }
    return res;
}

int main() {
    fast_io();

    int n, k;
    std::cin >> n >> k;

    std::vector<int> a(n);
    for (int i = 0; i < n; ++i) {
        std::cin >> a[i];
    }
    // 预处理：排序和去重
    std::sort(a.begin(), a.end());
    a.erase(std::unique(a.begin(), a.end()), a.end());
    
    // 找到最小价值，为我们的优化做准备
    int min_a = a[0];

    // 计算平移后的价值 b_i = a_i - min_a
    std::vector<int> b;
    for (int x : a) {
        b.push_back(x - min_a);
    }
    
    int max_b = 0;
    if (!b.empty()) {
        max_b = b.back();
    }
    
    // 构造初始多项式 p(x)，p[v] = 1 表示价值 v 存在
    std::vector<long long> p(max_b + 1, 0);
    for (int val : b) {
        p[val] = 1;
    }

    // 计算 p(x)^k
    std::vector<long long> res_poly = poly_pow(p, k);

    // 计算基础价值偏移量
    long long base_cost = (long long)k * min_a;
    bool first = true;
    // 遍历结果多项式，输出所有可能的总价值
    for (size_t i = 0; i < res_poly.size(); ++i) {
        if (res_poly[i] > 0) {
            if (!first) {
                std::cout << " ";
            }
            std::cout << base_cost + i;
            first = false;
        }
    }
    std::cout << std::endl;

    return 0;
}

复杂度分析，看看我们跑得有多快？

• 时间复杂度: O(M log M * log k) 的说。 这里的 M 是我们需要处理的多项式的最大长度，它约等于 k * (max(a_i) - min(a_i))。在最坏情况下，M 大约是 1000 * 1000 = 10^6。FFT 的长度 N 会是大于 M 的最小的 2 的幂。每一次多项式乘法的时间是 O(N log N)。我们用快速幂计算 k 次方，需要 O(log k) 次乘法。所以总时间就是 O(N log N * log k) 啦！对于这道题的数据范围，完全可以接受。

• 空间复杂度: O(M) 的说。 我们主要的内存开销在于存储多项式的系数向量。在 FFT 过程中，我们需要大小为 N 的向量来存储点值表示。所以空间复杂度和 N (也就是 M) 是一个数量级的。

知识点与总结，这次我们捕获了什么？

这次探险真是收获满满呀！我们成功捕获了以下这些闪闪发光的知识点：

1. 生成函数: 这是解决组合计数问题的超级大杀器！将组合问题转化为多项式代数问题，思路非常巧妙。以后遇到“求和的组合”这类问题，可以多往这个方向想一想哦。
2. FFT/NTT: 快速傅里叶变换是实现多项式快速乘法的核心算法。它是算法竞赛中一个非常重要的工具，尤其是在处理卷积问题时。这里的 NTT (数论变换) 是 FFT 在模意义下的版本，可以避免浮点数精度问题，不过这道题用 FFT 就足够啦。
3. 快速幂: 不仅仅能用于数值计算，它的思想可以推广到任何具有结合律的运算上，比如矩阵乘法和我们这里的多项式乘法。
4. 平移思想: 对数据进行整体平移（比如都减去最小值）是一种非常实用的优化技巧。它能简化问题，缩小数据范围，让计算更高效。

希望这次的题解能帮助你更好地理解 FFT 和生成函数的魅力！算法的世界就像一个巨大的游乐园，总有新奇又好玩的项目在等着我们。保持好奇心，继续探索吧！喵~ (ฅ'ω'ฅ)