0
\$\begingroup\$

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)
\$\endgroup\$
1
1
\$\begingroup\$

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, [], [])
\$\endgroup\$
1
0
\$\begingroup\$

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)
\$\endgroup\$
1
  • \$\begingroup\$ It's a bit O(n^2), though, isn't it? \$\endgroup\$ – Gareth Rees Nov 3 '13 at 22:27

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.