E - Cookies

E - Cookies

Problem Statement

There are $10^{100}$ cookies of each of $N$ types. The deliciousness per cookie of the $i$-th type is $A_i$.

You will choose a total of $K$ cookies from these. Two ways of choosing cookies are considered the same if and only if the multisets of types of chosen cookies match.

For each of the $\binom{N+K-1}{K}$ ways of choosing, consider the sum of deliciousness of the chosen cookies. Let $S_1,S_2,\dots$ be these sums in descending order, with duplicates included. Find $S_1,\dots,S_X$.

Constraints

• $1\leq N \leq 50$
• $1 \leq K \leq 10^5$
• $1 \leq X \leq \min\left(10^5, \binom{N+K-1}{K}\right)$
• $-10^9 \leq A_i \leq 10^9$
• All input values are integers.

---

Input

The input is given from Standard Input in the following format:

$N$ $K$ $X$
$A_1$ $\dots$ $A_N$

Output

Let $S_1,S_2,\dots$ be the possible sums of deliciousness of the chosen $K$ cookies in descending order, with duplicates included. Output $S_1,\dots,S_X$ in this order, separated by newlines.

---

Sample Input 1

2 4 3
20 10

Sample Output 1

80
70
60

There are $5$ ways to choose $4$ cookies: "$k$ cookies of type $1$ and $4-k$ cookies of type $2$" $(0\leq k \leq 4)$, where the sums of deliciousness of the chosen cookies are $80,70,60,50,40$.

---

Sample Input 2

3 100000 5
-1 -1 -1

Sample Output 2

-100000
-100000
-100000
-100000
-100000

Different ways of choosing cookies may result in the same sum of deliciousness.

---

Sample Input 3

9 14142 13
31 41 59 26 53 58 97 93 23

Sample Output 3

1371774
1371770
1371766
1371762
1371758
1371754
1371750
1371746
1371742
1371738
1371736
1371735
1371734

 

解题思路

  先说一下总的方案数 $\binom{N+K-1}{K}$ 是怎么来的。由于饼干之间没有差异，选择的方式只与其所属种类有关，因此可以看作是球盒模型。把饼干看作是小球，现在要将 $K$ 个相同的小球分配到 $N$ 个不同的盒子中，并允许某些盒子为空，那么该问题对应的方案数就是 $\binom{N+K-1}{K}$。

  接下来，介绍本题的解法。由于数组的排序对结果没有影响，因此我们可以假设 $a_1 \geq a_2 \geq \cdots \geq a_n$。定义一个长度为 $n$ 的序列 $c_i$ 来表示当前方案中选择了 $a_i$ 的数量。那么该方案对应的总得分就是 $s = \sum_{i=1}^{n}{c_i \cdot a_i}$。我们的目标是按得分从大到小依次列出所有可能的分数。可以使用类似于 bfs 的思路，从当前方案出发，逐步枚举得到下一个最大的分数及其对应的方案。并使用大根堆来维护当前已知但还未输出的候选方案，每次取出当前最大值，并从中扩展出下一个可能稍小的方案。

  容易知道选 $k$ 个 $a_1$ 时，能够得到最大得分 $k \cdot a_1$。那么如何从一个方案，扩展到分数严格小于等于当前，但尽可能大的相邻方案呢？

  考虑从当前方案 $(c_1,\dots,c_n)$ 出发，要得到一个不同的新方案，必须减少某种 $i$ 的数量 $1$ 个，同时增加某种 $j > i$ 的数量 $1$ 个，即执行一次 $c_i \gets c_i-1,\ c_j \gets c_j+1$ 操作。新方案的分数变化为 $a_j - a_i$，因为我们希望新方案的分数尽可能接近原分数，因此应选择最大的 $a_j$，即 $j = i+1$。选择变化量为 $1$ 也是同样的原因。因此从当前方案出发，最有希望成为次优的相邻方案，就是对每一个满足 $c_i > 0$ 的 $i$ 执行 $c_i \gets c_i-1,\ c_j \gets c_j+1$。

  还有一个问题是为什么需要枚举所有满足 $c_i > 0$ 的 $i$，而不是只选 $a_{i+1}-a_i$ 最大的那个 $i$？如果每次只走下降最少的一条边，就会退化成一条链，导致漏掉一些分数排名很靠前但不在该链上的方案。举个例子，有 $a=[11,10,1]$ 和 $k=2$，所有方案分数为 $\{22,21,20,12,11,2\}$。如果只走下降最少的路径，生成的方案序列为（元组的第一个元素是分数）：$$(22,2,0,0) \xrightarrow{c_1-1, c_2+1} (21,1,1,0) \xrightarrow{c_1-1, c_2+1} (20,0,2,0) \xrightarrow{c_2-1, c_3+1} (11,0,1,1)$$

  可以看到，方案 $(12, 1, 0, 1)$ 被忽略了。因此必须对所有可能的 $i$ 都进行扩展，才能保证不遗漏。

  综上所述，做法是先将数组 $a$ 按降序排列，使用 std::set（用于对方案去重）来维护一个包含 $n+1$ 元组 $(s, c_1, \dots, c_n)$ 的集合。每次从集合中选出分数最大的方案，并枚举所有的 $i$，将满足 $c_i > 0$ 的 $i$ 对应的方案 $(s - a_i + a_{i+1}, c_1, \dots, c_i - 1, c_{i+1}, \dots, c_n)$ 加入集合。通过这种方法，可以模拟选出前 $m$ 个最大分数。

  类似的题目还有 F - K-th Largest Triplet。

  AC 代码如下，时间复杂度为 $O\left( nm \log{nm} \right)$：

#include <bits/stdc++.h>
using namespace std;

typedef long long LL;

const int N = 55;

int a[N];

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, m, k;
    cin >> n >> m >> k;
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
    }
    sort(a + 1, a + n + 1, greater<int>());
    set<array<LL, 51>, greater<array<LL, 51>>> st;
    st.insert({{1ll * a[1] * m, m}});
    while (k--) {
        auto p = *st.begin();
        st.erase(st.begin());
        cout << p[0] << '\n';
        for (int i = 1; i < n; i++) {
            if (!p[i]) continue;
            p[0] += a[i + 1] - a[i], p[i]--, p[i + 1]++;
            st.insert(p);
            p[0] -= a[i + 1] - a[i], p[i]++, p[i + 1]--;
        }
    }
    
    return 0;
}

 

参考资料

  Editorial - AtCoder Beginner Contest 440：https://atcoder.jp/contests/abc440/editorial/15027