CODEWITHSUNDEEP
pythonrecursiontree
cyrus
cyrus

Reputation: 1356

How to recursively sum and store all child values in a tree

Given a tree, what's the easiest way to calculate the sum of all children at a given node?

Say a tree like this...enter image description here

The red values represent the sum of the node and its children.

Let's say the node structure look like this (an example):

class Node:
    def __init__(self, name):
        self.children = []
        self.weight = 100
        self.weight_plus_children = 295

How can I do this in an efficient, single pass (in Python)?

Thanks!

Upvotes: 2

Views: 5404

Answers (3)

lqhcpsgbl
lqhcpsgbl

Reputation: 3782

Just judge if a node is a leaf and add the sum to the weight, here is an example:

class Node:
    def __init__(self, name, weight, children):
        self.children = children
        self.weight = weight
        self.weight_plus_children = weight

    def get_all_weight(self):
        if self.children is None:
          return self.weight_plus_children
        else:
          for child in self.children:
            print "child.get_all_weight()", child.get_weigth_with_children()
            self.weight_plus_children += child.get_weigth_with_children()

        return self.weight_plus_children

    def get_weigth_with_children(self):
        return self.weight_plus_children

leaf1 = Node('C1', 58, None)
leaf2 = Node('C2', 7, None)
leaf3 = Node('C3', 10, None)
leaf4 = Node('C4', 20, None)

subroot = Node('B1', 50, [leaf1, leaf2])
subroot1 = Node('B2', 50, [leaf3, leaf4])

root = Node('A', 100, [subroot, subroot1])

print subroot.get_all_weight()
print
print subroot1.get_all_weight()
print
print root.get_all_weight()

OutPut:

F:\so>python test-tree.py
child.get_all_weight() 58
child.get_all_weight() 7
115

child.get_all_weight() 10
child.get_all_weight() 20
80

child.get_all_weight() 115
child.get_all_weight() 80
295

Upvotes: 3

Paul Lo
Paul Lo

Reputation: 6138

You can try the following recursive program:

def calculate_sum(cur):
    children_sum = 0
    for child in cur.children:
        children_sum += calculate_sum(child)
    return cur.weight + children_sum 

calculate_sum(root)  # pass in your root node

Upvotes: 0

jack
jack

Reputation: 2204

A function that adds its weight to the weights of its children, where the weights of its children are defined in terms of the function itself. Call it on the root node, and it should return the sum for the whole tree.

def treesum(node):
    return node.weight + sum(treesum(child) for child in node.children)

Upvotes: 0

Related Questions