[docs]class PriorityStack(object):
def __init__(self):
self.m_stack = [[]]
[docs] def pop(self):
return self.m_stack.pop()
[docs] def peek(self):
return self.m_stack[-1]
[docs] def push(self, m):
self.m_stack.append([m])
[docs] def push_last(self, m):
self.m_stack[-1].append(m)
[docs] def peek_last(self):
return self.m_stack[-1][-1]
[docs] def peek_level(self, level):
"""Find a whole level of nodes in the stack"""
for layer in self.m_stack:
if layer and layer[0][0] == level:
return [node for _, node in layer]
return []
[docs] def lineage(self):
"""Fetch the last element of each level of priorities. When the
stack is used to keep track of a tree, this list includes a list of
'parents', as the last element of each level is the parent being
processed."""
if self.m_stack[0]:
return list(reversed([els[-1][-1] for els in self.m_stack]))
else:
return []
[docs] def lineage_with_level(self):
if self.m_stack[0]:
return list(reversed([els[-1] for els in self.m_stack]))
else:
return []
[docs] def size(self):
return len(self.m_stack)
[docs] def unwind(self):
"""Combine nodes as needed while walking back up the stack. Intended
to be overridden, as how to combine elements depends on the element
type."""
raise NotImplementedError
[docs] def add(self, node_level, node):
""" Add a new node with level node_level to the stack. Unwind the stack
when necessary. Returns `self` for chaining """
last = self.peek()
element = (node_level, node)
if len(last) > 0 and node_level > last[0][0]:
self.push(element)
elif len(last) > 0 and node_level < last[0][0]:
while last[0][0] > node_level:
self.unwind()
last = self.peek()
self.push_last(element)
else:
self.push_last(element)
return self