# How to enumerate the internal nodes and leaves of a tree more elegantly?

Might there be a simpler way to write this code? I'd like to learn to write code more efficiently.

def leaves_and_internals(self):
"""(BTNode) -> ([object], [object])
Return a pair of lists: (list of all the items stored in the leaves of
the subtree rooted at this node, list of all the items stored in the
internal nodes of the subtree rooted at this node).
"""
leaves = []
internal = []

if self.left is None and self.right is None:
leaves.append(self.item)
else:
internal.append(self.item)
if self.left:
subleaf, subinternal = self.left.leaves_and_internals()
leaves.extend(subleaf)
internal.extend(subinternal)
if self.right:
subleaf, subinternal = self.right.leaves_and_internals()
leaves.extend(subleaf)
internal.extend(subinternal)

return (leaves, internal)

• You don't seem to have taken the comments on your previous question into account. Nov 3 '13 at 22:16

Here's a more efficient version that avoids copying list elements. Instead, it creates just two empty lists and appends to them during the traversal.

def leaves_and_internals(self):
"""(BTNode) -> ([object], [object])
Return a pair of lists: (list of all the items stored in the leaves of
the subtree rooted at this node, list of all the items stored in the
internal nodes of the subtree rooted at this node), according to a
pre-order traversal.
"""
def helper(self, leaves, internal):
if self.left is None and self.right is None:
leaves.append(self.item)
else:
internal.append(self.item)
for child in [self.left, self.right]:
if child is not None:
helper(child, leaves, internal)
return (leaves, internal)
return helper(self, [], [])


In the docstring, I would mention that the traversal is done using pre-order.

I suggest handling leaf nodes as a base case and returning early.

The left and right child nodes can be handled using common code.

def leaves_and_internals(self):
"""(BTNode) -> ([object], [object])
Return a pair of lists: (list of all the items stored in the leaves of
the subtree rooted at this node, list of all the items stored in the
internal nodes of the subtree rooted at this node), according to a
pre-order traversal.
"""

if self.left is None and self.right is None:
return ([self.item], [])

leaves, internal = [], [self.item]
for child in [self.left, self.right]:
if child is not None:
subleaf, subinternal = child.leaves_and_internals()
leaves.extend(subleaf)
internal.extend(subinternal)
return (leaves, internal)

• It's a bit O(n^2), though, isn't it? Nov 3 '13 at 22:27