numpy_heap.py

   1 import numpy
   2
   3 def bubble_up(heap, i, key):
   4   # call with key == tuple(heap[i, :])
   5   while i:
   6     parent = (i - 1) >> 1
   7     if key >= tuple(heap[parent, :]):
   8       break
   9     heap[i, :] = heap[parent, :]
  10     i = parent
  11   heap[i, :] = numpy.array(key, heap.dtype)
  12
  13 def bubble_down(heap, i, key, n):
  14   # call with key == tuple(heap[i, :])
  15   child = (i << 1) + 1
  16   while child < n:
  17     child_key = tuple(heap[child, :])
  18     child1 = child + 1
  19     if child1 < n:
  20       child1_key = tuple(heap[child1, :])
  21       if child1_key < child_key:
  22         child = child1
  23         child_key = child1_key
  24     if key <= child_key:
  25       break
  26     heap[i, :] = heap[child, :]
  27     i = child
  28     child = (i << 1) + 1
  29   heap[i, :] = numpy.array(key, heap.dtype)
  30
  31 def heapify(heap, n):
  32   for i in range((n - 1) >> 1, -1, -1):
  33     bubble_down(heap, i, tuple(heap[i, :]), n)