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中，那么so一定是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

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