Skip to content

Instantly share code, notes, and snippets.

@ijkilchenko
Last active June 16, 2016 04:19
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save ijkilchenko/9920bf05ef1a45420abb9f8f5adc2e7d to your computer and use it in GitHub Desktop.
Save ijkilchenko/9920bf05ef1a45420abb9f8f5adc2e7d to your computer and use it in GitHub Desktop.
Heap which supports insertion in O(log(n)) time and finding the max (the root node) in O(1). Also, the tree of the heap is always as short as possible -- what I call a balanced heap.
import random
class Node():
def __init__(self, data=None):
self.data = data
self.parent = None
self.left = None
self.right = None
self.load = 1 # Assumed to be a leaf when initialized.
class Heap():
def __init__(self):
self.root = Node()
def _bubble(self, node): # O(log(n))
# Let's recalculate the load (or the number of nodes below).
node.load = 1
if node.left is not None:
node.load += node.left.load
if node.right is not None:
node.load += node.right.load
# Now recalculate the parent's load until we reach the root node.
if node.parent is not None:
self._bubble(node.parent)
# TODO: Can left and right methods be further factored out?
def _attach_left(self, node, n_data):
# Here node.left is None so we place a new node with data n_data there.
n_node = Node(n_data)
n_node.parent = node
node.left = n_node
self._bubble(node.left)
def _attach_right(self, node, n_data):
# Here node.right is None so we place a new node with data n_data there.
n_node = Node(n_data)
n_node.parent = node
node.right = n_node
self._bubble(node.right)
def insert(self, n_data, node=None): # O(2*log(n))=O(log(n))
if node is None:
node = self.root # Ugh, why can't I access self in the signature, Python?
if node.data is None:
node.data = n_data
else:
# Node we are trying to make is less than the root.
if node.data >= n_data:
if node.data != n_data:
if node.left is None:
self._attach_left(node, n_data)
else:
self.insert(n_data, node.left)
else:
if node.right is None:
self._attach_right(node, n_data)
else:
self.insert(n_data, node.right)
# Node we are trying to make is greater than the root.
else:
# Swap the root with the new data and try to insert
# original root's data below.
temp_data = node.data
node.data = n_data
if node.left is not None and node.right is not None:
# When neither child is None, we decide which way
# to go based on the least loaded child (keeping heap balanced).
if node.left.load <= node.right.load:
self.insert(temp_data, node.left)
else:
self.insert(temp_data, node.right)
else:
if node.left is None:
self._attach_left(node, temp_data)
else:
self._attach_right(node, temp_data)
if __name__ == '__main__':
L = [1, 30, 3, 4, 1, 2, -1, -10, 40]
h = Heap()
for l in L:
h.insert(l)
print('L=%s' % L)
print('max element is %s' % h.root.data)
print('left child of root is %s' %h.root.left.data)
print('right child of root is %s' % h.root.right.data)
print('number of nodes of heap is %s' % h.root.load)
print()
L = [-20, 20, 100, 80, 1, 2, 3, 90]
h = Heap()
for l in L:
h.insert(l)
print('L=%s' % L)
print('max element is %s' % h.root.data)
print('left child of root is %s' %h.root.left.data)
print('right child of root is %s' % h.root.right.data)
print('number of nodes of heap is %s' % h.root.load)
print()
L = random.sample(list(range(1000)), 50)
h = Heap()
for l in L:
h.insert(l)
print('L=%s' % L)
print('max element is %s' % h.root.data)
assert max(L) == h.root.data
print('left child of root is %s' %h.root.left.data)
print('right child of root is %s' % h.root.right.data)
print('number of nodes of heap is %s' % h.root.load)
# NOTES:
# Children of the root are not always the second and third max elements (heap property).
# The node's load helps decide whether we attach left or right (to keep the heap balanced).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment