HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonMinor

Calculate the path sum of a binary tree

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
paththebinarycalculatesumtree

Problem

I want to calculate all sum of all possible paths from root to leaf. Each node has value of an integer number (decimal).

In the following example:

1
   / \
  2   4
 /    /\
3    5  6


Possible paths are *123, 145, and 146, which adds to 414). In this example, 12 is not a valid path since 2 is not a leaf node.

I did basic tests and it works, and I'm wondering if the code is logically correct, and how I can improve its performance.

class Node:

    def __init__(self, value, left, right):
        self.value = value
        self.left = left
        self.right = right

def sum(node, prefix):
    r1 = 0
    r2 = 0
    if not node:
        return
    if node and not node.left and not node.right:
        return prefix*10+node.value
    if node.left:
        r1 = sum(node.left, prefix*10 + node.value)
    if node.right:
        r2 = sum(node.right, prefix*10 + node.value)

    return r1+r2

if __name__ == "__main__":

    node3 = Node(3, None, None)
    node5 = Node(5, None, None)
    node6 = Node(6, None, None)
    node4 = Node(4, node5, node6)
    node2 = Node(2, node3, None)
    node1 = Node(1, node2, node4)

    print sum(node1, 0)

Solution

The code can be simplified a bit.

In the sequence:

if not node:
    return
if node and not node.left and not node.right:
    return prefix*10+node.value


...the if node part of the second condition is redundant. Since this is preceded by if not node: return, the second condition will never be reached unless if node evaluates to true.

For the if not node.left and not node.right, I'd apply DeMorgan's theorem, to get if node.left or node.right, and switch the controlled statement with the other.

By having it return 0 for an empty tree as @Zondo suggested, we can simplify a bit more: rather than checking whether node is there before summing it, we can just sum it, and if nothing is there, it returns 0, which doesn't change the sum.

With those and @Zondo's suggestions, we end up with something like this for the sum routine:

def sum(node, prefix = 0):

    if not node:
        return 0

    p = prefix * 10 + node.value

    if node.left or node.right:
        return sum(node.left, p) + sum(node.right, p)

    return p


Another possibility that might be worth considering would be to get a little closer to the single-entry, single-exit "ideal" that used to be advocated in structured-programming:

def sum(node, prefix = 0):

    if not node:
        return 0

    p = prefix * 10 + node.value

    if node.left or node.right:
        p = p + sum(node.left, p) + sum(node.right, p)

    return p


It would be pretty easy to transform this into a true single-entry, single-exit form, but I don't see it as a real improvement (and even what I've done above might be open to some question).

Code Snippets

if not node:
    return
if node and not node.left and not node.right:
    return prefix*10+node.value
def sum(node, prefix = 0):

    if not node:
        return 0

    p = prefix * 10 + node.value

    if node.left or node.right:
        return sum(node.left, p) + sum(node.right, p)

    return p
def sum(node, prefix = 0):

    if not node:
        return 0

    p = prefix * 10 + node.value

    if node.left or node.right:
        p = p + sum(node.left, p) + sum(node.right, p)

    return p

Context

StackExchange Code Review Q#134226, answer score: 5

Revisions (0)

No revisions yet.