Parsing Binary Tree Array in Python I got this from a practice problem set, and looking for feedback on my approach.


  Suppose you're given a binary tree represented as an array. Write a function that determines whether the left or right branch of the tree is larger. The size of each branch is the sum of the node values. The function should return the string "Right" if the right side is larger and "Left" if the left side is larger. If the tree has 0 nodes or if the size of the branches are equal, return the empty string.


My logic is that each "breadth" of the tree is 2^depth wide, and we start the work on the second "depth". I can use that to establish how to slice the array to get each "breadth" out of the array (ie, first round is arr[1:3], second is arr[3:7], etc). From there I sum the left/right slices of that slice.

I think my solution uses \$O(n)\$ time and \$O(1)\$ space and should handle unbalanced trees, but parsing trees as arrays is new to me, so wondering about edge cases I should be considering.

def solution(arr):
    breadth = 2
    floor = 1
    left = 0
    right = 0

    while floor &lt; len(arr):
        level = arr[floor:(floor+breadth)]
        pivot = breadth / 2

        left += sum(level[:pivot])
        right += sum(level[pivot:])

        floor += breadth
        breadth *=2

    if left &lt; right:
        return 'Right'
    elif right &lt; left:
        return 'Left'
    else:
        return ''

import unittest

class Test(unittest.TestCase):

    def test_function(self):
        self.assertEqual(solution([3,6,2,9,-1,10]), 'Left')
        self.assertEqual(solution([1,4,100,5]), 'Right')
        self.assertEqual(solution([1,10,5,1,0,6]), '')
        self.assertEqual(solution([]), '')
        self.assertEqual(solution([1]), '')

t = Test()
t.test_function()


One alternate I can think of is evaluating each value. Could this be better than my current solution?

for i in range(len(level)):
    if i &lt; pivot:
        left += level[i]
    else:
        right += level[i]
I think the solution you have is likely to be optimal (or close enough) in terms of runtime, and I think your approach of summing over the slice is much better than doing a for loop.  

The only downside of this approach is that it's a little hard to grok, and consequently difficult to extend or maintain, because the approach you're taking is so specific to the implementation of the tree and to the particular task you're doing.  In other words, this code would probably not be reusable for a different task that involved traversing the same tree, or for applying the same task to a different tree representation.  

If efficiency is of paramount importance, there's a place for tightly optimized code, but in that case I'd recommend having some explanatory comments to speed the reader's comprehension.  Something with a diagram of the tree layout, say:

# Our tree:
#                      arr[0]                      &lt;-  breadth 1, floor 0
#          arr[1]                  arr[2]          &lt;-  breadth 2, floor 1
#    arr[3]      arr[4]      arr[5]      arr[6]    &lt;-  breadth 4, floor 3

# Begin at the second level (top of the left and right subtrees).
breadth = 2  # number of nodes on the current level
floor = 1    # index of the first node on the current level
left = 0     # sum of the left subtree up to this point
right = 0    # sum of the right subtree up to this point


This is the diagram I built in my head when I was reading the code; having it in front of me already would have saved me a couple of minutes.  With that diagram and the explanations of the different values, it's a lot easier to read through the code and visualize what the slices correspond to.



An entirely different direction to go in for this problem would be to build an abstraction that lets you traverse the tree in a more tree-like way:

from typing import List, NewType, Optional

def solution(arr: List[int]) -&gt; str:
    """Takes a binary tree packed into an array root-first.
    Returns 'Right' or 'Left' to indicate which subtree is larger."""

    # Helpers to read and traverse nodes in the tree.
    # The internal representation of our Node is simply an array index, but
    # these helper functions expose derived "properties" as if it were an object.
    Node = NewType('Node', int)

    def make_node(index: int) -&gt; Optional[Node]:
        return Node(index) if index &lt; len(arr) else None

    def left_child(node: Node) -&gt; Optional[Node]:
        return make_node(node * 2 + 1)

    def right_child(node: Node) -&gt; Optional[Node]:
        return make_node(node * 2 + 2)

    def root_node() -&gt; Optional[Node]:
        return make_node(0)

    def value(node: Node) -&gt; int:
        return arr[node]

    # Recursive function to sum the subtree under a node.
    def sum_subtree(node: Optional[Node]) -&gt; int:
        if node is None:
            return 0
        return (
            value(node) 
            + sum_subtree(left_child(node)) 
            + sum_subtree(right_child(node))
        )

    # Get sums of left and right subtrees.
    root = root_node()
    if root is None:
        return ''
    left_sum = sum_subtree(left_child(root))
    right_sum = sum_subtree(right_child(root))

    if right_sum &gt; left_sum:
        return 'Right'
    elif left_sum &gt; right_sum:
        return 'Left'
    else:
        return ''


This approach breaks the logic into layers.  The foundational layer is the five functions that define the Node abstraction, which gives us a way to read the tree that's completely decoupled from the array.  The details of how the tree is stored in the array (e.g. the magic 2 * index + 1 logic) are completely contained inside these functions.  You could also implement this collection of functions as a class (storing a reference to the array as an instance attribute), which would make it easy to use outside this function and unit-test on its own; I've used a NewType wrapper because it makes it easy to see (and typecheck) the abstraction while still giving us the low overhead of a single int at runtime.

This tree abstraction lets us build a sum_subtree helper that recursively sums up everything under a node.  Even if the reader didn't understand the way the tree is packed into an array (or care to), they could read this function and understand how it's navigating the tree to arrive at a sum, and if they wanted to write a similar function they'd be able to do so without having to understand the array representation.

Finally, we have the top-level logic that gets the two top-level subtrees, compares their sums, and returns a result in the expected format.

This solution is a bit less efficient than yours (the recursive call is going to use log(N) space, and the fact that we're iterating through each node one by one will add a few more clock cycles), but in many real world applications "slow and obvious" is preferable over "fast and clever".  :)