例如，在机器学习中的自然语言处理中，通常使用波束搜索来预测序列中接下来要添加的对象并对其进行排序。波束搜索的关键部分是top-k评分指标，其基本原理是：给定一个长度为N的概率分数选择列表，返回N中得分最高的k个项目。这就像对列表进行排序然后取前几位的值那么简单。
参考波束搜索中的一个视觉示例https://www.researchgate.net/figure/A-partially-completed-beam-search-procedure-with-a-beam-width-of-5-for-an-example-input_fig2_317377611（在上面的例子中，k=5，且‘top’分数是最小值），在每次迭代中，每个节点从选择列表N中选择top k项目，产生k²个总潜在路径。从这些路径中，筛选出整体top k，这些成为下一轮迭代的节点。在前面的例子中，你可以看到每个时间步只显示了筛选后的节点。https://d2l.ai/_images/beam-search.svg全面扩展了k=2，N=5的情况。

想象一下，不是为每个分支/节点优化一个选择，而是必须选择多个值：当从一个节点探索时，你有一个维度为（N，q）的选择集，你想从中选择q个值，每列q中选一个。然后，要找到得分最高的选择集，你需要考虑这些列中值的组合。例如：对于一个选择矩阵N=5，q=4：

+---+--------+--------+--------+--------+
| N |   q0   |   q1   |   q2   |   q3   |
+---+--------+--------+--------+--------+
| 0 | 0.9763 | 0.0791 | 0.1530 | 0.5565 |
| 1 | 0.1560 | 0.1014 | 0.6932 | 0.7551 |
| 2 | 0.8142 | 0.9494 | 0.4582 | 0.4411 |
| 3 | 0.3807 | 0.2403 | 0.6897 | 0.7356 |
| 4 | 0.0156 | 0.9419 | 0.9568 | 0.2266 |
+---+--------+--------+--------+--------+

如果k=5，这个top-k函数应该返回以下结果：

1. 3.6376 = q₀[0] + q₁[2] + q₂[4] + q₃[1]
2. 3.6301 = q₀[0] + q₁[4] + q₂[4] + q₃[1]
3. 3.6181 = q₀[0] + q₁[2] + q₂[4] + q₃[3]
4. 3.6106 = q₀[0] + q₁[4] + q₂[4] + q₃[3]
5. 3.4755 = q₀[2] + q₁[2] + q₂[4] + q₃[1]

这些是使用每列中的一个值所能得到的最大可能总和。

对于任意N和q的解决方案，简单的方法是计算所有N^q的总和，然后对其排序，取前k个结果。第一步优化是排序每列，然后只计算每列中前k个值的总和组合，这样复杂度降低到k^q。

然而，考虑到这个寻找最高分的函数在波束搜索的每个时间步都需要调用k次，如果希望扩展到高k或高q，任何可能的加速都是至关重要的。我想到的最佳解决方案（简化为最小的例子，假设matrix是一个形状为（N，q）的numpy数组，并假设q为4）：

import numpy as npfrom itertools import combinationsclass Beamsearch():    def __init__(self, klen, q=4):        self.klen = klen        self.combis = []        for lens in range(klen):            self.combis.extend(list(self.partition(lens, q)))        self.width = q        self.wdth = list(range(q))    def partition(self, N, size):        n = N + size - 1        for splits in combinations(range(n), size - 1):            yield [s1 - s0 - 1 for s0, s1 in zip((-1,) + splits, splits + (n,))]    def getkmaxscores(self, matrix):        matrix_argsort = np.argsort(-matrix, axis=0)        sums = []        for comb in self.combis:            midxs = matrix_argsort[comb, self.wdth]            midxslist = midxs.tolist()            msum = (sum(matrix[midxs, self.wdth]),                    midxslist)            sums.append(msum)        sums.sort(reverse=True)        return sums[:self.klen]

这种方法为给定宽度q的整数p创建分区，对于整数0 ≤ p ≤ k，例如，对于q=4：

p0: [0, 0, 0, 0]p1: [0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]p2: [0, 0, 0, 2], [0, 0, 1, 1], [0, 0, 2, 0], [0, 1, 0, 1], [0, 1, 1, 0], [0, 2, 0, 0], [1, 0, 0, 1], [1, 0, 1, 0], [1, 1, 0, 0], [2, 0, 0, 0]

等等。

然后这些被用来索引排序后的输入矩阵，以选择每个组合进行求和。在q=4的情况下，p_i的长度遵循三角金字塔序列（https://oeis.org/A000292）：这将搜索空间减少到所有p_0…k的总和，这是二项式系数(k,4) = k(k-1)(k-2)(k-3)/24（https://oeis.org/A000332）。对于小k（对于k < 30，这是小于k³），这是一个巨大的改进，但仍然以k⁴的顺序增长。对于任意情况，是否存在复杂度<O(k^q)的解决方案？

回答：

这个问题在文献中被称为从X + Y中选择。经典参考是Frederickson和Johnson，他们给出了当X和Y排序时最优的O(k)时间算法。你的列未排序，F&J的算法相当复杂，所以让我简要介绍更简单的O(k log k)算法。

首先，对于X和Y，选择并排序前k个元素。初始化一个最大堆，其中元素(i, j)的优先级是X[i] + Y[j]。插入(0, 0)。重复以下步骤k次：弹出最大元素(i, j)并记录其优先级。插入(i, j+1)。如果j = 0，还要插入(i+1, 0)。这所有操作的时间复杂度为O(n + k log k)，其中n是列中元素的数量。

最后，让我们将问题简化为两列。如果有超过两列，例如X, Y, Z，那么我们可以从X + Y中选择前k个元素，然后从(X + Y) + Z中选择前k个元素。