如何找到具有第k大和的配对？

Question

给定两个排序后的数字数组，我们要找到第k大可能总和的一对数。（一对数是来自第一个数组的一个元素和来自第二个数组的一个元素）。例如，当输入的数组为

• [2, 3, 5, 8, 13]
• [4, 8, 12, 16]

时，有最大和的数对为

• 13 + 16 = 29
• 13 + 12 = 25
• 8 + 16 = 24
• 13 + 8 = 21
• 8 + 12 = 20

因此第4大的数对是(13, 8)。如何找到第k大可能的总和的数对呢？同时，最快的算法是什么？这些数组已经排好序，大小分别为M和N。

---

我已经知道了使用最大堆的O(KlogK)解决方案，该解决方案在这里给出。它也是谷歌面试中最喜欢的问题之一，并且他们要求O(k)的解决方案。我还读到过一种O(k)的解决方案，但我无法理解。能否有人用伪代码解释正确的解决方案？

P.S.请不要将这个链接作为答案/评论发布。它不包含答案。

- Spandan

---

你可以在这个PDF中找到一个线性时间算法："Selection in X + Y and matrices with sorted rows and columns"。 - Evgeny Kluev

你对问题的描述相当模糊。它是两个具有相同索引的元素的总和吗？这两个数组的元素数量是否相等？ - Nikolai Ruhe

@EvgenyKluev，PDF中描述的算法是**O(n)，而不是完全适用于O(k)**，它仅适用于相同长度的M和N。 - hs3180

@hs3180：是的，这个算法的时间复杂度为O(n)，比要求的O(k)要好。如果k<n，我们可以忽略所有下标大于k的数组元素（并使n=k）。如果n<k<n^2，我们可以获得更好的复杂度，即O(n)<O(k)。此外，如果M<N，我们可以总是将一些非常小的元素附加到最短的数组中（并使M=N）。 - Evgeny Kluev

可能是已排序矩阵上的选择算法的重复问题。 - David Eisenstat

Answer 1

我从一个简单但不是完全线性时间复杂度的算法开始。我们选择介于array1[0]+ array2 [0]和array1[N-1]+ array2 [N-1]之间的某个值。然后，我们确定有多少对总和大于此值，以及其中有多少对较小。这可以通过使用两个指针迭代数组来完成：当总和太大时，第一个数组的指针会递增，当总和太小时，第二个数组的指针会递减。通过为不同值重复执行此过程并使用二分搜索（或单侧二分搜索），我们可以在O（NlogR）时间内找到第K个最大的总和，其中N是最大数组的大小，R是array1[N-1]+ array2 [N-1]和array1[0]+ array2 [0]之间可能的值的数量。当数组元素由小常量限制时，此算法的时间复杂度仅为线性。

如果我们在二分搜索范围内的配对和数量从O(N²)降至O(N)时停止二分搜索，那么可以改进先前的算法。然后，我们用这些配对和填充辅助数组（可以使用略微修改的双指针算法完成）。接下来，我们使用快速选择算法在这个辅助数组中找到第K大的和。所有这些都不会改善最坏情况复杂度，因为我们仍然需要进行O(log R)次二分搜索步骤。如果我们保留此算法的快速选择部分，但使用比二分搜索更好的方法（以获得适当的值范围），会发生什么？

我们可以通过以下方法估计值的范围：从每个数组中获取每第二个元素，并尝试找到这些半数组的排名k/4的对和（使用相同的算法递归地进行）。显然，这应该给出所需值范围的一些近似值。实际上，这种技巧的稍微改进的变体给出了仅包含O(N)元素的范围。这在以下论文中得到证明："Selection in X + Y and matrices with sorted rows and columns" by A. Mirzaian and E. Arjomandi。该论文包含有关算法、证明、复杂性分析以及除Quickselect之外所有部分的伪代码的详细说明。如果需要线性最坏情况复杂度，则可以使用Median of medians算法增强Quickselect。

这个算法的复杂度为O(N)。如果其中一个数组比另一个数组短（M < N），我们可以假设这个较短的数组被扩展到大小N，其中包含一些非常小的元素，以便算法中的所有计算都使用最大数组的大小。实际上，我们不需要提取这些“添加”的元素对并将它们馈送到快速选择中，这使得算法稍微快一点，但不会改善渐近复杂度。

如果k < N，则可以忽略索引大于k的所有数组元素。在这种情况下，复杂度等于O(k)。如果N < k < N(N-1)，则我们只比OP请求的复杂度更好。如果k > N(N-1)，我们最好解决相反的问题：第k小的和。

我上传了简单的C++11实现到ideone。代码没有经过优化和彻底测试。我尽量使其接近链接论文中的伪代码。该实现使用std::nth_element，平均仅允许线性复杂度（而不是最坏情况）。

---

一种完全不同的方法来在线性时间内找到第K个和是基于优先队列（PQ）的。其中一个变化是将最大的一对插入到PQ中，然后重复删除PQ的顶部并插入最多两对（一个数组中下标递减，另一个数组中下标递减）。并采取一些措施防止插入重复的对。另一种变化是插入包含第一个数组中最大元素的所有可能对，然后重复删除PQ的顶部并插入第一个数组中下标递减且第二个数组中相同下标的对。在这种情况下，无需担心重复。

OP提到了使用最大堆实现PQ的O(K log K)解决方案。但在某些情况下（当数组元素是均匀分布的整数且仅需要平均线性复杂度而不是最坏情况下），我们可以使用O(1)时间优先队列，例如，在这篇论文中所述："A Complexity O(1) Priority Queue for Event Driven Molecular Dynamics Simulations" by Gerald Paul。这使得期望时间复杂度为O(K)。

这种方法的优点是能够按排序顺序提供前K个元素。缺点是数组元素类型选择有限，算法更加复杂和缓慢，渐进复杂度更差：O(K) > O(N)。

- Evgeny Kluev

---

你能否发布一个可运行的示例？这样太复杂了，难以理解。 - Spandan

我有一个简单得多的解决方案。我邀请您找出其中的缺陷。 - Marcin

你应该接受这个答案。这个算法并不简单，而这可能是你在互联网上能找到的最好的解释。 - Santtu Keskinen

@EvgenyKluev：赏金已发放，实至名归 :) - Spandan

这是一个经过深入研究的答案，但并没有回答简单可信的O(k)解决方案的要求。我还不会放弃，预计几个小时内会有一篇文章。 - clwhisk

Answer 2

[2, 3, 5, 8, 13]
[4, 8, 12, 16]

合并这两个数组，并记录排序后数组中的索引。以下是索引数组的样子（从1开始而不是0）：

[1, 2, 4, 6, 8] [3, 5, 7, 9]

现在从末尾开始创建元组，将元组中的元素相加，并选择第k大的和。

- Rohit Kumar Gupta

Answer 3

public static List<List<Integer>> optimization(int[] nums1, int[] nums2, int k) {
            // 2 * O(n log(n))
            Arrays.sort(nums1);
            Arrays.sort(nums2);
            List<List<Integer>> results = new ArrayList<>(k);
            int endIndex = 0;
            // Find the number whose square is the first one bigger than k
            for (int i = 1; i <= k; i++) {
                if (i * i >= k) {
                    endIndex = i;
                    break;
                }
            }
            // The following Iteration provides at most endIndex^2 elements, and both arrays are in ascending order,
            // so k smallest pairs must can be found in this iteration. To flatten the nested loop, refer
            // 'https://dev59.com/X2s05IYBdhLWcg3wFOC8'
            for (int i = 0; i < endIndex * endIndex; i++) {
                int m = i / endIndex;
                int n = i % endIndex;
                List<Integer> item = new ArrayList<>(2);
                item.add(nums1[m]);
                item.add(nums2[n]);
                results.add(item);
            }
            results.sort(Comparator.comparing(pair->pair.get(0) + pair.get(1)));
            return results.stream().limit(k).collect(Collectors.toList());
        }

消除O(n^2)的关键：

1. 避免两个数组的笛卡尔积（或类似于交叉连接的操作），也就是展开嵌套循环。

2. 缩小对两个数组的迭代范围。

所以：

1. 对两个数组进行排序（根据Java文档，Arrays.sort提供O(n log(n))的性能）

2. 将迭代范围限制在刚好足够支持查找k个最小对的大小。

- vince

Answer 4

简而言之：如果在每次迭代中向前和向后查看，您可以从结束处（即最高点）开始并以O(K)时间往回工作。

尽管这种方法的基本思想是正确的，但下面的代码目前还不完全正确（请参见注释）。

---

让我们看一下：首先，这些数组是排序过的。因此，如果数组是长度为M和N的a和b，您已经按照顺序排列它们，那么最大的项分别在槽M和N中，则最大的对始终为a[M]+b[N]。

现在，第二个最大的对是什么？它可能会有{a[M]，b[N]}之一（因为两者都不行，那只是最大的对），并且至少有{a[M-1]，b[N-1]}之一。但是，我们还知道如果选择a[M-1]+b[N-1]，我们可以通过选择来自同一列表的更高数字使其中一个操作数更大，因此它将恰好具有来自最后一列和倒数第二列的数字。

考虑以下两个数组：a = [1, 2, 53]; b = [66, 67, 68]。我们的最高匹配是 53+68。如果我们失去其中较小的那一个，我们的匹配是68+2;如果我们失去较大的那一个，它是53+67。因此，我们必须向前看来决定我们的下一个匹配。最简单的前瞻策略是计算两个可能的匹配的总和。每次转换都将始终花费两个加法和两个比较（三个因为我们需要处理总和相等的情况）；让我们称之为成本Q。

起初，我想重复K-1次。但有个问题：下一个最大的匹配可能实际上是我们可以从{{a[M]，b[N]}，{a[M-1]，b[N-1]}}中制作的另一对。因此，我们还需要向后查找。

因此，让我们编写代码（使用Python，应该兼容2/3）：

def kth(a,b,k):
    M = len(a)
    N = len(b)
    if k > M*N:
       raise ValueError("There are only %s possible pairs; you asked for the %sth largest, which is impossible" % M*N,k)
    (ia,ib) = M-1,N-1 #0 based arrays
    # we need this for lookback
    nottakenindices = (0,0) # could be any value
    nottakensum = float('-inf')
    for i in range(k-1):
        optionone = a[ia]+b[ib-1]
        optiontwo = a[ia-1]+b[ib]
        biggest = max((optionone,optiontwo))
        #first deal with look behind
        if nottakensum > biggest:
           if optionone == biggest:
               newnottakenindices = (ia,ib-1)
           else: newnottakenindices = (ia-1,ib)
           ia,ib = nottakenindices
           nottakensum = biggest
           nottakenindices = newnottakenindices
        #deal with case where indices hit 0
        elif ia <= 0 and ib <= 0:
             ia = ib = 0
        elif ia <= 0:
            ib-=1
            ia = 0
            nottakensum = float('-inf')
        elif ib <= 0:
            ia-=1
            ib = 0
            nottakensum = float('-inf')
        #lookahead cases
        elif optionone > optiontwo: 
           #then choose the first option as our next pair
           nottakensum,nottakenindices = optiontwo,(ia-1,ib)
           ib-=1
        elif optionone < optiontwo: # choose the second
           nottakensum,nottakenindices = optionone,(ia,ib-1)
           ia-=1
        #next two cases apply if options are equal
        elif a[ia] > b[ib]:# drop the smallest
           nottakensum,nottakenindices = optiontwo,(ia-1,ib)
           ib-=1
        else: # might be equal or not - we can choose arbitrarily if equal
           nottakensum,nottakenindices = optionone,(ia,ib-1)
           ia-=1
        #+2 - one for zero-based, one for skipping the 1st largest 
        data = (i+2,a[ia],b[ib],a[ia]+b[ib],ia,ib)
        narrative = "%sth largest pair is %s+%s=%s, with indices (%s,%s)" % data
        print (narrative) #this will work in both versions of python
        if ia <= 0 and ib <= 0:
           raise ValueError("Both arrays exhausted before Kth (%sth) pair reached"%data[0])
    return data, narrative

对于没有Python的人，这里有一个ideone：http://ideone.com/tfm2MA

最坏情况下，每次迭代中有5个比较，K-1次迭代，这意味着这是一个O(K)算法。

现在，可能可以利用值之间的差异信息来优化一下，但这已经实现了目标。

---

这是一个参考实现（不是 O(K) 的，但除非存在一些对于和相等的情况有特殊要求的情况，否则总是有效的）：

import itertools
def refkth(a,b,k):
    (rightia,righta),(rightib,rightb) = sorted(itertools.product(enumerate(a),enumerate(b)), key=lamba((ia,ea),(ib,eb):ea+eb)[k-1]
    data = k,righta,rightb,righta+rightb,rightia,rightib
    narrative = "%sth largest pair is %s+%s=%s, with indices (%s,%s)" % data
    print (narrative) #this will work in both versions of python
    return data, narrative

这个程序计算两个数组的笛卡尔积（即所有可能的对），按照它们的和排序，并取第k个元素。 enumerate 函数为每个项目添加了索引。

- Marcin

---

我看到这段代码有两个缺陷。(1)你从未进行任何边界检查。例如，在ia-=1之后，你很容易让ia等于-1。Python方便地从末尾重新开始遍历相同的数组。但这显然不是预期的。(2)如果nottakensum > biggest，你将一个“向前看”的位置放到nottakenindices中，并完全忽略另一个位置，在其他两个选项中，你只是覆盖nottakenindices而没有使用它的先前值。因此，在所有3种情况下，你都会丢失一些应该稍后考虑的位置（在正确的实现中）。 - Evgeny Kluev

@EvgenyKluev（1）说得好（2）这些是关于我的代码所做的正确观察，但它们并没有揭示错误。您能否提供一个这是错误的情况的例子？ - Marcin

这段代码定义了两个数组：a=[1,2,3] 和 b=[10,20]。在找到三个正确的和（23、22、21）之后，由于越界而没有更多的“lookahead”位置，并且nottakenindices只有(0,0)位置。因此，这段代码找到的第四个和是11，而13和12则被忽略了。 - Evgeny Kluev

@EvgenyKluev 谢谢。我会考虑一下。 - Marcin

仅查看1个元素是不够的。考虑两个数组{5,50,70,80,90,100}：第一组小于50 + 50的最小值是90 + 5。在我看来，O(k)似乎非常乐观。 - vgru

@Groo 确实如此。我需要重新考虑这个问题。 - Marcin

Answer 5

编辑：这个不行。我保留答案，因为显然不只有我会有这样的想法；请参见下面的讨论。反例是 x = (2, 3, 6)，y = (1, 4, 5)，k=3，其中算法给出 7 (3+4)，而不是 8 (3+5)。

---

假设有两个按降序排列的数组 x 和 y，我们想要构建第 K 大的和。

变量包括：在第一个数组中的索引 i（元素为 x[i]）、在第二个数组中的索引 j（元素为 y[j]）以及和的“顺序” k（k 在 1..K 范围内），意思是 S(k)=x[i]+y[j] 将是满足条件的第 k 大的和（这是循环不变式）。

从 (i, j) 等于 (0, 0) 开始：显然，S(1) = x[0]+y[0]。

对于 k 从 1 到 K-1，执行以下操作：

• 如果 x[i+1]+ y[j] > x[i] + y[j+1]，那么 i := i+1（而j不变）；否则 j:=j+1

为了验证它的工作原理，考虑你有 S(k) = x[i] + y[j]。那么，S(k+1) 是最大的和，它小于（或等于）S(k)，并且至少有一个元素（i 或 j）发生了变化。很容易看出，i 或 j 中的一个应该发生变化。

如果 i 发生变化，你可以构造一个比 S(k) 更小的更大的和，方法是设置 i=i+1，因为 x 是递减的，并且所有的 x[i'] + y[j]（其中 i' < i）都大于 S(k)。对于 j 也是一样的，这表明 S(k+1) 要么是 x[i+1] + y[j]，要么是 x[i] + y[j+1]。

因此，在循环结束时，你找到了第 K 大的和。

- Nicolas Grebille

---

这与我几天前发布的有缺陷的解决方案基本相同。问题的一部分是它不能生成所有可能的和，因此根据逻辑推理不会总是找到第K大的和。考虑数组{2, 10}和{1, 2}，这里有4个可能的和{3, 4, 11, 12}。您的算法只会生成其中的3个{3, 4, 12}。最初，我将我的错误答案发布出来，并讨论了其中的问题，以帮助其他人，比如你自己，不犯同样的错误。然而，在吸引了负评之后，我决定将其删除。 - NealB

我不太理解。第一步是(10, 2) -> 12。第二步，我比较了(2, 2)->4和(10, 1)->11，通过减少每个索引来选择(10,1)->11。您为什么会说该算法无法生成11？ - Nicolas Grebille

只要数组排序，数组的大小并不重要。该算法生成确切地 K 个和，这些和是数组中最大的 K 个和（因为它们都不同，并且 没有一个 非生成和可以通过构造变得更大）。请注意，原帖要求 O(K) 的算法，很明显不能尝试 N*M 种可能的和！ - Nicolas Grebille

当我选择K为N*M时会发生什么？ - NealB

@NealB 我明白了，谢谢。我会留下答案并解释一下，以防有人会有相同（错误）的想法。 - Nicolas Grebille

Answer 6

如果最后两个解分别在(a1, b1), (a2, b2)，那么在我的看法中，只有四个可能的解(a1-1, b1)、(a1, b1-1)、(a2-1, b2)、(a2, b2-1)。这种直觉可能是错误的。每个坐标最多有四个候选项，并且下一个最高的是在16个对中出现(a in {a1,a2,a1-1,a2-1}, b in {b1,b2,b1-1,b2-1})。这是O(k)。

(不，它不是，我仍然不确定是否可能。)

- clwhisk

---

在我看来，仅仅向后查找一个元素是不够的。考虑以下两个数组：{5, 50, 70, 80, 90, 100}。第一对小于 50+50 的数是 90+5。 - vgru

正确。我认为关键是只需查看 k 个元素，总共不超过 O(k)。在这个例子中，5+90 是第 26 高的一对，找到它需要比大多数其他内容更多的工作。我在捕捉这种推理线路方面遇到了麻烦，并写下了这个不太成功的尝试！ - clwhisk

你说的“往前看O(k)个元素”是什么意思？我不明白为什么选择k会影响排序对。你是不是想说，在每次迭代中都回溯到最大的元素? - vgru

一次回退一个元素，直到回退超过45个元素。那么需要回退的总次数（在找到所有k个最高和的操作计数之和上）是否受k的倍数限制？ - clwhisk

Answer 7

另一个问题中的最大堆算法简单、快速且正确。不要否定它。它也解释得非常清楚。https://dev59.com/a2435IYBdhLWcg3wyzQJ#5212618

也许没有任何O(k)算法，但是O(k log k)几乎一样快。

- Colonel Panic

---

实际上，我认为堆的复杂度是O(k log n)——其中n是两个向量中较短的那个的长度。由于对于任何给定的向量对来说，n都是常数，因此实际上是O(k)。参见：后来相关问题的答案。 - user3793679