给出两个列表，找到前K个产品的有效方法

Question

Given two lists of equal length N , I want to find the K largest products that can be made by multiplying an element from each list. For example, if

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

the result is [90, 72, 70] or [9*10, 9*8, 7*10] , found by

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

Is there a more efficient algorithm that doesn't involve multiplying out all N^2 pairs?

Answer 1

As already noted, the first step is to sort both lists A and B in descending order (or just the K largest of both lists). Then, all the max K products will sit in a roughly triangular area in the top-left corner, the max product being A[0]*B[0] . In other words, if A[i]*B[j] is in the top K, then so must be both A[i-1]*B[j] and A[i]*B[j-1] (assuming i, j > 0 ).

Thus, you can start in the top-left corner and then use a Heap to expand both the "lower" and the "right" neighbor of the current element and put those onto the heap, too, until you have all the K elements you need. Or start with all the K largest elements of A paired with the largest from B already on the heap and only expand in one direction.

Example in Python, using the heapq module, but the same will work in almost any other language. Note that we are adding negative products to the heap as the heap will be sorted smallest-first.

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

Example:

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

The complexity should be about O(NlogN + KlogK) for sorting A and B and then about K iterations with heap-operations in the loop. Each cell in the triangular "target" region will only be expanded once from its left neighbor, and cells added to the heap but not used are also limited to K (one in each "row"), giving a maximum of 2*K elements inspected.

Answer 2

Practical solution:

Find largest K elements from list A and K largest elements from list B by using partial_sort (this is a well-known modification of quick sort, and I am sure python has the same in its library). Largest products formed by these new lists are also the largest products of the original lists. Then use max-heap (priority queue) to find K largest products from new lists.

Answer 3

If we would find out K max values from both the lists, we would have the max K products from both the lists.

I would suggest two approaches to find out K max values:

1. If K <<< N ( K in 10s and N in millions )
   Here you have couple of options.
   • You can use selection algorithm K times for both the lists. That would take O(N*K)
   • K iterations of either Selection Sort or Bubble Sort . You would have K max values at either at the beginning or at the end of the array depending on the type of implementation. Even that would be O(N*K)

Note that because K <<< N you can say that O(N*K) is almost O(N)

2. K can be as same as N
   • In this case, The best bet would be to just sort both the lists using Merge Sort or Quick Sort . That would be O(N*lgN)