python - 给定两个列表查找前 K 产品的有效方法

Question

给定两个长度相等的列表 N，我想找到 K 最大的乘积，这些乘积可以通过将每个列表中的一个元素相乘得到。例如，如果

> A = [7, 9, 4, 1, 6]
> B = [8, 1, 3, 10, 7]
> K = 3

结果为[90, 72, 70] 或[9*10, 9*8, 7*10]，由

> sorted([x*y for x in A for y in B], reverse=True)[:K]
[90, 72, 70]

有没有更高效的算法，不需要将所有 N^2 对相乘？

Answer 1

如前所述，第一步是按降序对列表 A 和 B 进行排序(或者仅对两个列表中的 K 个最大的列表进行排序)。然后，所有最大 K 个产品将位于左上角的大致三角形区域中，最大产品为 A[0]*B[0]。换句话说，如果 A[i]*B[j] 在前 K 个中，那么 A[i-1]*B[j] 也一定是和 A[i]*B[j-1](假设 i, j > 0)。

因此，您可以从左上角开始，然后使用 Heap扩展当前元素的“低”和“右”邻居，并将它们也放入堆中，直到拥有所需的所有 K 元素。或者从 A 的所有 K 个最大元素与 B 中已经在堆上的最大元素配对开始，并且只向一个方向扩展。

Python 示例，使用 heapq模块，但同样适用于几乎任何其他语言。请注意，我们将负数产品添加到堆中，因为堆将按最小优先排序。

def top_k_prod(A, B, k):
    A = heapq.nlargest(k, A)
    B = heapq.nlargest(k, B)
    result = []
    heap = [(-A[i] * B[0], i, 0) for i in range(len(A))]
    while heap and len(result) < k:
        p, a, b = heapq.heappop(heap)
        result.append(-p)
        if b < len(B)-1:
            heapq.heappush(heap, (-A[a] * B[b+1], a, b+1))
    return result

例子:

import random
A = [random.randint(0, 100) for _ in range(100)]
B = [random.randint(0, 100) for _ in range(100)]
K = 20
result = top_k_prod(A, B, K)
test = sorted([x*y for x in A for y in B], reverse=True)[:K]
print(result)
# [9900, 9702, 9603, 9600, 9504, 9408, 9405, 9405, 9400, 9400, 9312, 9306, 9300, 9216, 9212, 9212, 9207, 9200, 9120, 9120]
print(result == test)
# True

复杂度应该是 O(NlogN + KlogK) 用于排序 A 和 B 然后大约 K 循环中的堆操作迭代。三角形“目标”区域中的每个单元格只会从其左邻域扩展一次，添加到堆但未使用的单元格也限制为 K(每个“行”中一个)，最多可提供 2*K 个元素检查。