CODEWITHSUNDEEP
pythonalgorithmrecursionbinary-tree
Dominic Farolino
Dominic Farolino

Reputation: 1383

Binary Tree Traversal Sum Of Each Depth

I am looking for an algorithm or modification of an efficient way to get the sum of run through of a tree depth, for example:

                Z
               / \
              /   \
             /     \
            /       \
           X         Y
          / \      /  \
         /   \    /    \
        A     B  C      D

The number of final leaves is four so us have four final sums and they would be this respectively.

[Z+X+A] [Z+X+B] [Z+Y+C] [Z+Y+D]

If someone could guide me in the right direction into getting the sums of all possible depths, that would be great.

This will be done in python with fairly large trees.

Upvotes: 2

Views: 1652

Answers (2)

twinlakes
twinlakes

Reputation: 10208

Here is what you are looking for. In this example, trees are stored as dicts with a "value" and a "children" keys, and "children" maps to a list.

def triesum(t):
    if not t['children']:
        return [t['value']]
    return [t['value'] + n for c in t['children'] for n in triesum(c)]

Example

trie = {'value': 5, 'children': [
    {'value': 7, 'children': [
        {'value': 8, 'children': []},
        {'value': 2, 'children': []}
    ]},
    {'value': 4, 'children': [
        {'value': 3, 'children': []},
        {'value': 6, 'children': []}
    ]}
]}

print sorted(triesum(trie)) == sorted([5 + 7 + 8, 5 + 7 + 2, 5 + 4 + 3, 5 + 4 + 6])
    # prints True

Upvotes: 0

Hesham Attia
Hesham Attia

Reputation: 977

You can recurse over the nodes of the tree keeping the sum from the root down to this point. When you reach a leaf node, you return the current sum in a list of one element. In the internal nodes, you concatenate the lists returned from children.

Sample Code:

class Node:
    def __init__(self, value, children):
        self.value = value
        self.children = children

def tree_sums(root, current_sum):
    current_sum += root.value
    if len(root.children) == 0:
        return [current_sum]
    subtree_sums = []
    for child in root.children:
        subtree_sums += tree_sums(child, current_sum)
    return subtree_sums

tree = Node(1, [Node(2, []), Node(3, [])])
assert tree_sums(tree, 0) == [3, 4]

Upvotes: 1

Related Questions