詳解Python中heapq模塊的用法

更新時(shí)間：2016年06月28日 17:39:47 作者：cangmean

Python中的heapq模塊提供了一種堆隊(duì)列heapq類型,這樣實(shí)現(xiàn)堆排序等算法便相當(dāng)方便,這里我們就來(lái)詳解Python中heapq模塊的用法,需要的朋友可以參考下

heapq 模塊提供了堆算法。heapq是一種子節(jié)點(diǎn)和父節(jié)點(diǎn)排序的樹形數(shù)據(jù)結(jié)構(gòu)。這個(gè)模塊提供heap[k] <= heap[2*k+1] and heap[k] <= heap[2*k+2]。為了比較不存在的元素被人為是無(wú)限大的。heap最小的元素總是[0]。

打印 heapq 類型

import math 
import random
from cStringIO import StringIO

def show_tree(tree, total_width=36, fill=' '):
   output = StringIO()
   last_row = -1
   for i, n in enumerate(tree):
     if i:
       row = int(math.floor(math.log(i+1, 2)))
     else:
       row = 0
     if row != last_row:
       output.write('\n')
     columns = 2**row
     col_width = int(math.floor((total_width * 1.0) / columns))
     output.write(str(n).center(col_width, fill))
     last_row = row
   print output.getvalue()
   print '-' * total_width
   print 
   return

data = random.sample(range(1,8), 7)
print 'data: ', data
show_tree(data)

打印結(jié)果

data: [3, 2, 6, 5, 4, 7, 1]

     3           
  2      6      
5    4  7     1   
-------------------------
heapq.heappush(heap, item)

push一個(gè)元素到heap里, 修改上面的代碼

heap = []
data = random.sample(range(1,8), 7)
print 'data: ', data

for i in data:
  print 'add %3d:' % i
  heapq.heappush(heap, i)
  show_tree(heap)

打印結(jié)果

data: [6, 1, 5, 4, 3, 7, 2]
add  6:
         6         
 ------------------------------------
add  1:
      1 
   6         
------------------------------------
add  5:
      1 
   6       5       
------------------------------------
add  4:
        1 
    4       5       
  6
------------------------------------
add  3:
        1 
    3       5       
  6    4
------------------------------------
add  7:
        1 
    3        5       
  6    4    7
------------------------------------
add  2:
        1 
    3        2       
  6    4    7    5
------------------------------------

根據(jù)結(jié)果可以了解，子節(jié)點(diǎn)的元素大于父節(jié)點(diǎn)元素。而兄弟節(jié)點(diǎn)則不會(huì)排序。

heapq.heapify(list)

將list類型轉(zhuǎn)化為heap, 在線性時(shí)間內(nèi), 重新排列列表。

print 'data: ', data
heapq.heapify(data)
print 'data: ', data

show_tree(data)

打印結(jié)果

data: [2, 7, 4, 3, 6, 5, 1]
data: [1, 3, 2, 7, 6, 5, 4]

      1         
   3         2     
7    6    5    4  
------------------------------------
heapq.heappop(heap)

刪除并返回堆中最小的元素, 通過(guò)heapify() 和heappop()來(lái)排序。

data = random.sample(range(1, 8), 7)
print 'data: ', data
heapq.heapify(data)
show_tree(data)

heap = []
while data:
  i = heapq.heappop(data)
  print 'pop %3d:' % i
  show_tree(data)
  heap.append(i)
print 'heap: ', heap

打印結(jié)果

data: [4, 1, 3, 7, 5, 6, 2]

         1
    4         2
  7    5    6    3
------------------------------------

pop  1:
         2
    4         3
  7    5    6
------------------------------------
pop  2:
         3
    4         6
  7    5
------------------------------------
pop  3:
         4
    5         6
  7
------------------------------------
pop  4:
         5
    7         6
------------------------------------
pop  5:
         6
    7
------------------------------------
pop  6:
        7
------------------------------------
pop  7:

------------------------------------
heap: [1, 2, 3, 4, 5, 6, 7]

可以看到已排好序的heap。

heapq.heapreplace(iterable, n)

刪除現(xiàn)有元素并將其替換為一個(gè)新值。

data = random.sample(range(1, 8), 7)
print 'data: ', data
heapq.heapify(data)
show_tree(data)

for n in [8, 9, 10]:
  smallest = heapq.heapreplace(data, n)
  print 'replace %2d with %2d:' % (smallest, n)
  show_tree(data)

打印結(jié)果

data: [7, 5, 4, 2, 6, 3, 1]

         1
    2         3
  5    6    7    4
------------------------------------

replace 1 with 8:

         2
    5         3
  8    6    7    4
------------------------------------

replace 2 with 9:

         3
    5         4
  8    6    7    9
------------------------------------

replace 3 with 10:

         4
    5         7
  8    6    10    9
------------------------------------

heapq.nlargest(n, iterable) 和 heapq.nsmallest(n, iterable)

返回列表中的n個(gè)最大值和最小值

data = range(1,6)
l = heapq.nlargest(3, data)
print l     # [5, 4, 3]

s = heapq.nsmallest(3, data)
print s     # [1, 2, 3]

PS：一個(gè)計(jì)算題
構(gòu)建元素個(gè)數(shù)為 K=5 的最小堆代碼實(shí)例:

#!/usr/bin/env python 
# -*- encoding: utf-8 -*- 
# Author: kentzhan 
# 
 
import heapq 
import random 
 
heap = [] 
heapq.heapify(heap) 
for i in range(15): 
 item = random.randint(10, 100) 
 print "comeing ", item, 
 if len(heap) >= 5: 
  top_item = heap[0] # smallest in heap 
  if top_item < item: # min heap 
   top_item = heapq.heappop(heap) 
   print "pop", top_item, 
   heapq.heappush(heap, item) 
   print "push", item, 
 else: 
  heapq.heappush(heap, item) 
  print "push", item, 
 pass 
 print heap 
pass 
print heap 
 
print "sort" 
heap.sort() 
 
print heap

結(jié)果:

2016628172708102.png (550×363)