def _balance(self, root):
# balancer will always produce a new balanced set
balancedSet = PersistentBalancedSet(iterator=self._iterator)
rightStack = Stack()
# last default values for first iteration
lastBalanceFactor = 0
lastNode = None
lastParent = None
# current default values for first iteration
parent = None # for root parent will be None
node = root.__root
isUnbalanced = False
while node:
balanceFactor = getBalanceFactor(node)
if balanceFactor < -1 or balanceFactor > 1:
isUnbalanced = True
if isUnbalanced and balanceFactor in (1, 0, -1):
if lastBalanceFactor < -1:
self._rightRotation(lastNode, lastParent, balancedSet)
else:
self._leftRotation(lastNode, lastParent, balancedSet)
return balancedSet
if not balancedSet.__root:
node = copy(node)
object.__setattr__(balancedSet, '_PersistentBalancedSet__root', node)
# WARNING: change order of execution of below section on your own risk.
if node.right:
rnode = copy(node.right)
node.right = rnode
rightStack.push((rnode, node,)) # a tuple of right node and its parent, preserving right elements
if node.left:
lnode = copy(node.left)
node.left = lnode
# useful for next iteration
lastNode = node
lastParent = parent
lastBalanceFactor = balanceFactor
# this is to set next node in iteration as per preorder traversal rules
if node.left:
parent = node
node = node.left
elif not rightStack.isEmpty():
element = rightStack.pop()
node = element[0]
parent = element[1]
else:
node = None
return balancedSet
class NodeInorderIterator(object):
def __init__(self, node):
self._node = node
self._accumulator = Stack()
def __iter__(self):
return self
def __next__(self):
if self._accumulator.isEmpty() and not self._node:
raise StopIteration()
if self._node:
current = self._node
while current:
self._accumulator.push(current)
current = current.left
node = self._accumulator.pop()
self._node = node.right
return node
class BinaryTreePreorderIterator(object):
def __init__(self, node):
self._node = node
self._accumulator = Stack()
def __iter__(self):
return self
def __next__(self):
if self._node:
if self._node.right:
self._accumulator.push(self._node.right)
node = self._node
self._node = self._node.left
else:
if self._accumulator.isEmpty() and not self._node:
raise StopIteration()
node = self._accumulator.pop()
if node.right:
self._accumulator.push(node.right)
self._node = node.left
return node.value