found some execution time; also ditch way to complicated heap implementation and just use the balancing binary tree
This commit is contained in:
Stefan Harmuth
parent
9d373596a1
commit
cdac29dee7
125
tools/trees.py
125
tools/trees.py
@ -32,7 +32,7 @@ class TreeNode:
|
||||
def update_node(node: TreeNode):
|
||||
left_depth = node.left.height if node.left is not None else -1
|
||||
right_depth = node.right.height if node.right is not None else -1
|
||||
node.height = 1 + max(left_depth, right_depth)
|
||||
node.height = 1 + (left_depth if left_depth > right_depth else right_depth)
|
||||
node.balance_factor = right_depth - left_depth
|
||||
|
||||
|
||||
@ -60,15 +60,17 @@ class BinarySearchTree:
|
||||
|
||||
if obj < node.value:
|
||||
node.left = self._insert(node.left, node, obj)
|
||||
else:
|
||||
elif obj > node.value:
|
||||
node.right = self._insert(node.right, node, obj)
|
||||
else:
|
||||
raise ValueError("obj already present in tree: %s" % obj)
|
||||
|
||||
update_node(node)
|
||||
return self._balance(node)
|
||||
|
||||
def add(self, obj: Any):
|
||||
if obj is None or obj in self:
|
||||
raise ValueError("obj is None or already present in tree")
|
||||
if obj is None:
|
||||
return
|
||||
|
||||
self.root = self._insert(self.root, self.root, obj)
|
||||
self.node_count += 1
|
||||
@ -94,6 +96,8 @@ class BinarySearchTree:
|
||||
node.parent.left = None
|
||||
else:
|
||||
node.parent.right = None
|
||||
else:
|
||||
self.root = None
|
||||
elif node.left is not None and node.right is not None: # both subtrees present
|
||||
d_node = node.left
|
||||
while d_node.right is not None:
|
||||
@ -106,6 +110,8 @@ class BinarySearchTree:
|
||||
node.parent.left = node.right
|
||||
else:
|
||||
node.parent.right = node.right
|
||||
else:
|
||||
self.root = node.right
|
||||
node.right.parent = node.parent
|
||||
else: # only a subtree on the left
|
||||
if node.parent is not None:
|
||||
@ -113,10 +119,13 @@ class BinarySearchTree:
|
||||
node.parent.left = node.left
|
||||
else:
|
||||
node.parent.right = node.left
|
||||
else:
|
||||
self.root = node.left
|
||||
node.left.parent = node.parent
|
||||
|
||||
update_node(root_node)
|
||||
self._balance(root_node)
|
||||
self.node_count -= 1
|
||||
return
|
||||
|
||||
raise ValueError("obj not in tree:", obj)
|
||||
@ -187,101 +196,39 @@ class BinarySearchTree:
|
||||
|
||||
|
||||
class Heap(BinarySearchTree):
|
||||
def _find_left(self) -> TreeNode:
|
||||
heap = Queue()
|
||||
heap.put(self.root)
|
||||
while c_node := heap.get():
|
||||
if c_node.left is None or c_node.right is None:
|
||||
return c_node
|
||||
else:
|
||||
heap.put(c_node.left)
|
||||
heap.put(c_node.right)
|
||||
|
||||
def _find_right(self) -> TreeNode:
|
||||
heap = [self.root]
|
||||
while c_node := heap.pop():
|
||||
if c_node.left is None and c_node.right is None:
|
||||
return c_node
|
||||
else:
|
||||
if c_node.left is not None:
|
||||
heap.append(c_node.left)
|
||||
if c_node.right is not None:
|
||||
heap.append(c_node.right)
|
||||
|
||||
def _sort_up(self, node: TreeNode):
|
||||
pass
|
||||
|
||||
def _heapify(self):
|
||||
pass
|
||||
|
||||
def empty(self) -> bool:
|
||||
def empty(self):
|
||||
return self.root is None
|
||||
|
||||
def add(self, obj: Any):
|
||||
node = TreeNode(obj)
|
||||
if self.root is None:
|
||||
self.root = node
|
||||
else:
|
||||
t_node = self._find_left()
|
||||
if t_node.left is None:
|
||||
t_node.left = node
|
||||
else:
|
||||
t_node.right = node
|
||||
node.parent = t_node
|
||||
self._sort_up(node)
|
||||
|
||||
def pop(self) -> Any:
|
||||
def popMin(self):
|
||||
if self.root is None:
|
||||
raise IndexError("pop from empty heap")
|
||||
|
||||
ret = self.root.value
|
||||
if self.root.left is None and self.root.right is None:
|
||||
self.root = None
|
||||
else:
|
||||
d_node = self._find_right()
|
||||
self.root.value = d_node.value
|
||||
if d_node.parent.left == d_node:
|
||||
d_node.parent.left = None
|
||||
else:
|
||||
d_node.parent.right = None
|
||||
self._heapify()
|
||||
c_node = self.root
|
||||
while c_node.left is not None:
|
||||
c_node = c_node.left
|
||||
|
||||
ret = c_node.value
|
||||
self.remove(ret, c_node.parent)
|
||||
return ret
|
||||
|
||||
def popMax(self):
|
||||
if self.root is None:
|
||||
raise IndexError("pop from empty heap")
|
||||
|
||||
c_node = self.root
|
||||
while c_node.right is not None:
|
||||
c_node = c_node.right
|
||||
|
||||
ret = c_node.value
|
||||
self.remove(ret, c_node.parent)
|
||||
return ret
|
||||
|
||||
|
||||
class MinHeap(Heap):
|
||||
def _sort_up(self, node: TreeNode):
|
||||
while node.parent is not None and node.value < node.parent.value:
|
||||
node.value, node.parent.value = node.parent.value, node.value
|
||||
node = node.parent
|
||||
|
||||
def _heapify(self):
|
||||
node = self.root
|
||||
while node is not None:
|
||||
if node.left and node.left.value < node.value and (not node.right or node.right.value > node.left.value):
|
||||
node.left.value, node.value = node.value, node.left.value
|
||||
node = node.left
|
||||
elif node.right and node.right.value < node.value and (not node.left or node.left.value > node.right.value):
|
||||
node.right.value, node.value = node.value, node.right.value
|
||||
node = node.right
|
||||
else:
|
||||
break
|
||||
def pop(self):
|
||||
return self.popMin()
|
||||
|
||||
|
||||
class MaxHeap(Heap):
|
||||
def _sort_up(self, node: TreeNode):
|
||||
while node.parent is not None and node.value > node.parent.value:
|
||||
node.value, node.parent.value = node.parent.value, node.value
|
||||
node = node.parent
|
||||
|
||||
def _heapify(self):
|
||||
node = self.root
|
||||
while node is not None:
|
||||
if node.left and node.left.value > node.value and (not node.right or node.right.value < node.left.value):
|
||||
node.left.value, node.value = node.value, node.left.value
|
||||
node = node.left
|
||||
elif node.right and node.right.value > node.value and (not node.left or node.left.value < node.right.value):
|
||||
node.right.value, node.value = node.value, node.right.value
|
||||
node = node.right
|
||||
else:
|
||||
break
|
||||
def pop(self):
|
||||
return self.popMax()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user