82 lines
1.8 KiB
Python
82 lines
1.8 KiB
Python
|
# node of the tree
|
||
|
|
class TreeNode:
|
||
|
|
def __init__(self):
|
||
|
|
self.item = None # Can be anything (a string currently)
|
||
|
|
# Only one parent can exist
|
||
|
|
self.parent = None
|
||
|
|
self.is_sample = False
|
||
|
|
self.children = []
|
||
|
|
|
||
|
|
# creates a child node
|
||
|
|
def create_leaf(self):
|
||
|
|
node = TreeNode()
|
||
|
|
node.parent = self
|
||
|
|
self.children.append(node)
|
||
|
|
return node
|
||
|
|
|
||
|
|
# returns true if this node is an "leaf" node, i.e:
|
||
|
|
# it has no children to continue with
|
||
|
|
def is_leaf(self):
|
||
|
|
if len(self.children) == 0:
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
def create_child_leaf(self, item):
|
||
|
|
node = self.create_leaf()
|
||
|
|
node.item = item
|
||
|
|
return node
|
||
|
|
|
||
|
|
def walk_children(self, fn):
|
||
|
|
fn(self)
|
||
|
|
# walk children nodes if this node isn't an leaf node
|
||
|
|
if not self.is_leaf():
|
||
|
|
for node in self.children:
|
||
|
|
node.walk_children(fn)
|
||
|
|
|
||
|
|
# find a single child. Returns None if no child with name ecists
|
||
|
|
def find_child(self, name):
|
||
|
|
if self.item == name:
|
||
|
|
return self
|
||
|
|
for node in self.children:
|
||
|
|
if node.item == name:
|
||
|
|
return node
|
||
|
|
return None
|
||
|
|
|
||
|
|
def get_parent_count(self):
|
||
|
|
# Walk the parents to figure out current tree depth
|
||
|
|
# and how many
|
||
|
|
parent = self.parent
|
||
|
|
parent_count = 0
|
||
|
|
while parent is not None:
|
||
|
|
parent_count += 1
|
||
|
|
parent = parent.parent
|
||
|
|
return parent_count
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
# def walk_children(self, fn):
|
||
|
|
# fn(self)
|
||
|
|
# for node in self.children:
|
||
|
|
# fn(node)
|
||
|
|
# for other_node in node.children:
|
||
|
|
# fn(other_node)
|
||
|
|
|
||
|
|
class Tree:
|
||
|
|
def __init__(self):
|
||
|
|
# create a root node
|
||
|
|
self.root = TreeNode()
|
||
|
|
#self.root.item = '(root)'
|
||
|
|
|
||
|
|
def walk(self, fn):
|
||
|
|
self.root.walk_children(fn)
|
||
|
|
#for node in self.root.children:
|
||
|
|
# node.walk_children(fn)
|
||
|
|
#fn(node)
|
||
|
|
#for other_node in node.children:
|
||
|
|
# fn(other_node)
|
||
|
|
|
||
|
|
# creates a leaf in the root node and populates it
|
||
|
|
def create_leaf(self, item):
|
||
|
|
leaf_node = self.root.create_child()
|
||
|
|
leaf_node.item = item
|
||
|
|
return leaf_node
|