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

Question

没能用谷歌搜索到这个问题的名称，希望这个问题能对社区有所贡献。

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

 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 (但它涉及创建一个完整的结果数组，而在我的例子中，我只需要前几个值)

附言我添加了 python 标签，因为它是一个数学问题，但我很乐意使用任何语言解决方案，或者只是解释，或者维基百科的链接...

Answer 1

你可以使用堆:

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)]