Last active
June 16, 2016 04:19
-
-
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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