从两个排序数组中取前 10 个元素的最小总和组合

Take first 10 element smallest sum combinations from two sorted arrays

没能google这个问题的名称,希望这个问题能对社区有所贡献。

假设我们有两个排序的数字数组,例如:

 2  8
12 18
45 35
85 48
87 49
97 59

并且我们希望首先从两个数组中有效地获取 k (10) 个数字的最小总和组合。在我们的例子中是:

 2 +  8 = 10 
 2 + 18 = 20
12 +  8 = 20
12 + 18 = 30
 2 + 35 = 37
12 + 35 = 47
 2 + 48 = 50
 2 + 49 = 51
45 +  8 = 53
12 + 48 = 60

什么是正确的方法?我草拟了一个简单的实现(由@sanyash 改进),但它没有利用数组已排序的事实,而且问题在线性时间内感觉可行...

def smallest_product(k, arr1, arr2):
    product_iter = itertools.product(
        itertools.islice(arr1, k),
        itertools.islice(arr2, k),
    )
    product_sorted = sorted(product_iter, key=sum)
    product_sliced = itertools.islice(product_sorted, k);
    return list(product_sliced)

print(smallest_product(10, 
    [ 2, 12, 45, 85, 87, 98], 
    [ 8, 18, 35, 48, 49, 59]))

类似问题:efficient sorted Cartesian product of 2 sorted array of integers(但它涉及创建一个完整的结果数组,而在我的例子中,我只需要前几个值)

P.S。我添加了 python 标签,因为它是一个数学问题,但我会很高兴用任何语言解决方案或只是解释,或者 link 到维基百科...

首先将两个数组截断为 len k。之后使用您之前的实现。难度为O(k^2):

import itertools

def smallest_product(k, arr1, arr2):
    product_iter = itertools.product(
        itertools.islice(arr1, k),
        itertools.islice(arr2, k),
    )
    product_sorted = sorted(product_iter, key=sum)[:k]
    return list(product_sorted)


print(smallest_product(
    10, 
    [ 2, 12, 45, 85, 87, 98], 
    [ 8, 18, 35, 48, 49, 59])
)

假设我们使用两个数组创建一个 table:

for arr in [[i + j for j in arr2] for i in arr1]: print(arr)

我们得到这样的输出:

[10, 20, 37, 50, 51, 61]
[20, 30, 47, 60, 61, 71]
[53, 63, 80, 93, 94, 104]
[93, 103, 120, 133, 134, 144]
[95, 105, 122, 135, 136, 146]
[106, 116, 133, 146, 147, 157]

请注意,在此矩阵中,matrix[i][j] == arr1[i] + arr2[j]。所以我们可以计算O(1)中矩阵任意位置的元素值。请注意,这是一个 排序矩阵 ,其中所有行和列都是单调递增的,我们试图在其中找到 k 最小的元素。

在这个阶段,O(KlogN) 堆方法变得相当简单。取第一行并将其变成最小堆。每次弹出最小的元素并将其添加到结果中。每次弹出时,都会将下一行对应列中的元素添加到堆中。重复 k 次,您将得到 k 最小的总和。

这与情况不完全相关,但是确实存在鞍背搜索的变体,可以让您在 O(N) 中找到排序矩阵中的 kth 最小元素,而不是 O(KlogN) 与上述方法一样。可能有一种方法可以将 this paper 中采用的方法修改为 O(K),但对于这种情况,这很可能是矫枉过正。

你可以使用堆:

import heapq


def smallest_product(k, a, b):
    k = min(k, len(a) * len(b))
    l = [(a[0] + b[0], 0, 0)]
    heapq.heapify(l)

    seen = set()
    for _ in range(k):
        s, i, j = heapq.heappop(l)

        if i + 1 < len(a) and (i + 1, j) not in seen:
            heapq.heappush(l, (a[i + 1] + b[j], i + 1, j))
            seen.add((i + 1, j))
        if j + 1 < len(b) and (i, j + 1) not in seen:
            heapq.heappush(l, (a[i] + b[j + 1], i, j + 1))
            seen.add((i, j + 1))
        yield (a[i], b[j])

result = list(smallest_product(10, [ 2, 12, 45, 85, 87, 98], [ 8, 18, 35, 48, 49, 59]))

print(result)

输出

[(2, 8), (2, 18), (12, 8), (12, 18), (2, 35), (12, 35), (2, 48), (2, 49), (45, 8), (12, 48)]

以上代码是 here 中代码的 Python 翻译。此方法的时间复杂度为 O(k*log k).

输出 (对于k = 11)

[(2, 8), (2, 18), (12, 8), (12, 18), (2, 35), (12, 35), (2, 48), (2, 49), (45, 8), (12, 48), (2, 59)]

这个问题可以分三步解决。

  1. 构造一个长度为m的已排序对的可迭代列表,其中m = min(len(list1), k)
  2. m-way merging (see this 应用于可迭代对象以获得排序对的单个可迭代对象,使用每对的总和作为键;
  3. 从 iterable 中取出前 k 个元素。

m 路合并有不同的算法。以下是基于堆的实现。复杂度为 O(k*logm).

from itertools import islice
from heapq import merge

def smallest_pairs(k, list1, list2):
    pairs = map(lambda x:((x, y) for y in list2), islice(list1, k))
    return list(islice(merge(*pairs, key=sum), k))

print(smallest_pairs(10, 
      [ 2, 12, 45, 85, 87, 98],
      [ 8, 18, 35, 48, 49, 59]))