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] <- breadth 1, floor 0
# arr[1] arr[2] <- breadth 2, floor 1
# arr[3] arr[4] arr[5] arr[6] <- 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]) -> 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) -> Optional[Node]:
return Node(index) if index < len(arr) else None
def left_child(node: Node) -> Optional[Node]:
return make_node(node * 2 + 1)
def right_child(node: Node) -> Optional[Node]:
return make_node(node * 2 + 2)
def root_node() -> Optional[Node]:
return make_node(0)
def value(node: Node) -> int:
return arr[node]
# Recursive function to sum the subtree under a node.
def sum_subtree(node: Optional[Node]) -> 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 > left_sum:
return 'Right'
elif left_sum > 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". :)