Calculate the path sum of a binary tree

Question

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)

Answer 1

What you really need is comments. Why are you multiplying prefix by 10? Let's see, if sum() is first called with ... I need to figure it out! Your code should have comments that clearly explain what is going on.

---

When I first glanced through your function, I noticed you calling sum() with two arguments. sum() takes only one argument (an iterable) unless you also have an initial argument. Oh, oh, you are calling your own sum function. It is a bad idea to shadow built-ins. If you actually wanted to use the built-in sum function, you might get some unexpected behavior that could be difficult to debug. Also, future readers (which might include you) could be confused, as I was.

---

You should use default arguments. If you want to create a leaf node, it is much easier to use Node(3) than Node(3, None, None). It is also nice to say Node(8, right=node2) than Node(8, None, node2). To implement, simply put =None after the left and right parameters when you define __init__(). Also, sum() would benefit from a default value of 0 for prefix. The person using the function might wonder, "Why do I always need to provide this argument? It is always the same." Making it a default argument means that it can be called with just sum(node1) instead of sum(node1, 0). If the user really wants to find the sum from the middle, he can still provide that argument. Starting from the middle, however, will probably always be something that only your function uses.

---

Your function is meant to return a number. Certain parts of the function depend on other calls returning a number. Therefore, I would change if not node: return to if not node: return 0

---

You use prefix*10 + node.value three times. It might be a good idea to define a variable for that. While it will be calculated a maximum of only two times per function call, it's something, and it's easy to fix.

Answer 2

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).