I'm building a tree with arbitrary structure. Leaves of the tree have a particular attribute that no other nodes have. I'm trying to come up with a way to traverse the tree to find a particular leaf and keep a running sum of the attributes of the leaves I've visited. I need the cumulative sum of the leaves I've visited that are not the leaf I'm searching for, as well as the attribute of the leaf I am trying to find. I need it to search left-to-right, top-to-bottom.
An example tree might look something like this:
My Node class is:
class Node:
def __init__(self, name, children=[], **kwargs):
self.__dict__.update(kwargs)
self.name = name
self.children = children
def insert(self, child):
self.children.append(child)
If I built the example tree as
l1 = Node("leaf1", attr=1)
l2 = Node("leaf2", attr=2)
n2 = Node("node2", children=[l1, l2])
l3 = Node("leaf3", attr=3)
l4 = Node("leaf4", attr=4)
n1 = Node("node1", children=[n2, l3, l4])
l6 = Node("leaf6", attr=6)
n4 = Node("node4", children=[l6])
l5 = Node("leaf5", attr=5)
l7 = Node("leaf7", attr=7)
n3 = Node("node3", children=[l5, n4, l7])
root = Node("root", children=[n1, n3])
and I wanted to search_node(root, "leaf4"), I'd expect a result of (6, 4) - the first element of the tuple is the sum of all leaves visited that were not leaf 4, and the second element is the attribute attached to leaf 4. search_node(root, "leaf1") would result in (0, 1), as there were no leaves visited that were not leaf 1.
Here's the code that I have so far:
def search_node(node, search_name):
cum_sum = 0
if node.children:
for child in node.children:
if child.name == search_name:
return (cum_sum, child.attr)
subtree_attr_sum, node_attr = search_node(child, search_name)
if node_attr is not None:
return (subtree_attr_sum, node_attr)
else:
cum_sum += subtree_attr_sum
return (cum_sum, None)
else:
return (node.attr, None)
It seems to work on the left half of the tree, but not on the right half. Admittedly, this code is ugly and could probably be reduced.