CODEWITHSUNDEEP
pythonrecursionbinary-tree
ofeks
ofeks

Reputation: 107

Sum of binary tree leaves' values

I wrote this code and when I use print I see that I get the leaves. However, the final return from the function is None and not the sum of the leaves, which is supposed to be 7 in this example. I'd be happy to know whats wrong here. Thank you !

class Node:
    def __init__(self, val=None):
        self.left = None
        self.right = None
        self.val = val


def sum_leafs(tree):
    if tree is None:
        return 0

    if tree.right and tree.left:
        sum_leafs(tree.right)
        sum_leafs(tree.left)

    elif tree.right or tree.left:
        if tree.right:
            sum_leafs(tree.right)
        elif tree.left:
            sum_leafs(tree.left)

    elif tree.right is None and tree.left is None:
        return sum_leafs(tree.left) + 1


node = Node(10)
node.right = Node(2)
node.left = Node(11)
node.left.right = Node(5)

print(sum_leafs(node))

Upvotes: 1

Views: 2047

Answers (4)

Mulan
Mulan

Reputation: 135217

node

First I'm going to update your Node interface so that it's possible to set left and right branches when creating nodes -

class Node:
  def __init__(self, val=None, left=None, right=None):
    self.left = left
    self.right = right
    self.val = val

This allows us to create tress more ergonomically, such as -

t = Node(10, Node(11, None, Node(5)), Node(2))

traverse

Now we write a generic traverse procedure. This allows us to separate 1) the traversal of our tree from 2) the intended operation we want to perform on each tree element -

def traverse(tree):
  if tree is None:
    return
  else:
    yield tree.val
    yield from traverse(tree.left)
    yield from traverse(tree.right)

Now the need for sum_leafs disappears. We have decoupled traversal logic from summing logic. We can calculate the sum of leafs with a simple combination of sum and traverse -

print(sum(traverse(t)))
# 28

don't repeat yourself

Or, instead of summing the values, we could write a search function to find the first value that passes a predicate -

def search(test, tree):
  for val in traverse(tree):
    if test(val):
      return val

print(search(lambda x: x < 10, t))
# 5

print(search(lambda x: x > 99, t))
# None

Or, we could simply collect each value into a list -

print(list(traverse(t)))
# [ 10, 11, 5, 2 ]

As you can see, removing the traversal logic from each function that depends on our tree can be a huge help.

without generators

If you don't like generators, you can write the eager version of traverse which always returns a list. The difference now is there is no way to partially traverse the tree. Note the similarities this program shares with the generator version -

def traverse(t):
  if t is None:
    return [] # <-- empty
  else:
    return \
      [ t.val
      , *traverse(t.left)  # <-- yield from
      , *traverse(t.right) # <-- yield from
      ]

print(traverse(t))
# [ 10, 11, 5, 2 ]

Upvotes: 0

Nir Alfasi
Nir Alfasi

Reputation: 53525

You forgot to add + when you sum the branches (left/right) and also you forgot to access val which is the most crucial thing for the whole thing to work.

Further, the logic can be simplified:

def sum_leafs(tree):
    if tree is None:
        return 0

    if not tree.right and not tree.left:
        return tree.val

    return sum_leafs(tree.right) + sum_leafs(tree.left)

Upvotes: 1

Akaisteph7
Akaisteph7

Reputation: 6486

You are not properly returning the calculated leaf sums. Try this:

class Node:
    def __init__(self, val=None):
        self.left = None
        self.right = None
        self.val = val


def sum_leafs(tree):
    if tree is None:
        return 0

    elif tree.right and tree.left:
        return sum_leafs(tree.right) + sum_leafs(tree.left)

    elif tree.right or tree.left:
        if tree.right:
            return sum_leafs(tree.right)
        elif tree.left:
            return sum_leafs(tree.left)

    elif tree.right is None and tree.left is None:
        return tree.val

node = Node(10)
node.right = Node(2)
node.left = Node(11)
node.left.right = Node(5)

print(sum_leafs(node))

7

Upvotes: 0

Jmonsky
Jmonsky

Reputation: 1519

You are not adding the sums together or returning them. This can also be done with a method in the class:

class Node:
    def __init__(self, val=None):
        self.left = None
        self.right = None
        self.val = val

    def sum(self):
        s = 0
        if self.left is not None:
            s += self.left.sum()
        if self.right is not None:
            s += self.right.sum()
        return self.val + s


node = Node(10)
node.right = Node(2)
node.left = Node(11)
node.left.right = Node(5)

print(node.sum())

returns:

28

Upvotes: 0

Related Questions