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

时间:2015-06-12 10:45:40

标签: python

我的目标是:

  1. 将元素分成两个列表;和
  2. 找到k最小的对。
  3. 目前我可以将两个列表中的每个元素与下面的代码配对:

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

    我的输出是:

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

    假设我设置k = 3,它应该返回三个最小的对:(1, 2)(1, 4)(3, 2)。我该怎么做?

1 个答案:

答案 0 :(得分:0)

您可以heapq.nsmallest使用sum作为其关键功能:

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

或者@jonrsharpe在评论中说你可以使用sorted

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