在两个列表中找到k个最小的对

Question

My goal is to:

1. pair the elements in two lists; and
2. find the k smallest pairs.

At the moment I can pair every element in both lists with the code below:

import itertools  

a = [1, 3, 5, 7]
b = [2, 4, 6, 8]
c = list(itertools.product(a, b))

print c

And my output is:

[(1, 2), (1, 4), (1, 6), (1, 8), (3, 2), (3, 4), (3, 6), (3, 8), (5, 2), (5, 4), (5, 6), (5, 8), (7, 2), (7, 4), (7, 6), (7, 8)]

Suppose I set k = 3 , it should return me the three smallest pairs: (1, 2) , (1, 4) and (3, 2) . How do I do that?

Answer 1

You can use heapq.nsmallest with sum as its key function :

>>> import heapq
>>> heapq.nsmallest(3,c,key=sum)
[(1, 2), (1, 4), (3, 2)]

Or as @jonrsharpe said in comment you can use sorted :

sorted(c, key=sum)[:k]