如何在Python中维护堆中的字典？

Question

我有一个字典，内容如下：

{'abc':100,'xyz':200,'def':250 .............}

这个字典的键是某个实体的名称，而值是这个实体的数量。我需要从这个字典中返回前10个元素。

我可以用堆来实现这个功能，但我不太确定怎么处理值和键的对应关系，因为有些值可能是相同的。

有没有其他的数据结构可以做到这一点呢？

Answer 1

如果你想获取前10个元素，假设这个数字在第二个位置：

from operator import itemgetter

topten = sorted(mydict.items(), key=itemgetter(1), reverse = True)[0:10]

如果你想按数值排序，然后再按键排序，只需要把它改成 key=itemgetter(1,0)。

至于数据结构，使用堆（heap）会比较合适。你只需要把它们保持为元组（tuple），然后比较数字部分。

Answer 2

使用堆是一种最佳解决方案，时间复杂度是：O(nlogk)。这里的n是堆的长度，而k是10。

关于键的映射，有个小技巧就是我们可以创建一个新的类来比较键，并定义一些特殊的方法 __lt__() 和 __gt__()，这些方法可以重写小于（<）和大于（>）运算符。

import heapq
class CompareWord:
  def __init__(self , word , value):
    self.word = word
    self.value = value

  def __lt__(self, other):   #To override > operator
    return self.value < other.value

  def __gt__(self , other):  #To override < operator
    return self.value > other.value

  def getWord(self):
    return self.word

def findKGreaterValues(compare_dict , k):
  min_heap = []
  for word in compare_dict:
      heapq.heappush(min_heap , CompareWord(word ,compare_dict[word] ))
      if(len(min_heap) > k):
          heapq.heappop(min_heap)   
  answer = []
  for compare_word_obj in min_heap:
      answer.append(compare_word_obj.getWord())

  return answer

Answer 3

使用 heapq 的话，你可能想这样做：

heap = [(-value, key) for key,value in the_dict.items()]
largest = heapq.nsmallest(10, heap)
largest = [(key, -value) for value, key in largest]

需要注意的是，因为 heapq 只实现了一个最小堆，所以最好把值反转一下，这样大的值就变成小的了。

对于小规模的堆，这种方法会比较慢，比如：

>>> import random
>>> import itertools as it
>>> def key_generator():
...     characters = [chr(random.randint(65, 90)) for x in range(100)]
...     for i in it.count():
...             yield ''.join(random.sample(characters, 3))
... 
>>> the_dict = dict((key, random.randint(-500, 500)) for key, _ in zip(key_generator(), range(3000)))
>>> def with_heapq(the_dict):
...     items = [(-value, key) for key, value in the_dict.items()]
...     smallest = heapq.nsmallest(10, items)
...     return [-value for value, key in smallest]
... 
>>> def with_sorted(the_dict):
...     return sorted(the_dict.items(), key=(lambda x: x[1]), reverse=True)[:10]
... 
>>> import timeit
>>> timeit.timeit('with_heapq(the_dict)', 'from __main__ import the_dict, with_heapq', number=1000)
0.9220538139343262
>>> timeit.timeit('with_sorted(the_dict)', 'from __main__ import the_dict, with_sorted', number=1000)
1.2792410850524902

当有3000个值时，它的速度只比 sorted 版本稍快，而 sorted 的时间复杂度是 O(nlogn)，而不是 O(n + mlogn)。如果我们把字典的大小增加到10000，heapq 的版本就会更快：

>>> timeit.timeit('with_heapq(the_dict)', 'from __main__ import the_dict, with_heapq', number=1000)
2.436316967010498
>>> timeit.timeit('with_sorted(the_dict)', 'from __main__ import the_dict, with_sorted', number=1000)
3.585728168487549

运行的机器也可能会影响时间。你最好测试一下哪种方法在你的情况下效果最好。如果效率不是特别重要，我建议使用 sorted 版本，因为它更简单。