'Size of subtree in Python

Almost every online tutorial I see on the subject when it comes to finding the size of a subtree involves calling a recursive function on each child's subtree.

The problem with this in Python is that it overflows if you recurse past a few hundred levels, so if I theoretically had a long, linear tree, it would fail.

Is there a better way to handle this? Do I need to use a stack instead?

pythonalgorithm


Solution 1:[1]

The stack would be the easiest non-recursive way of getting the size of the subtree (count of nodes under the given node, including the current node)

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

def subtree_size(root):
    visited = 0
    if not root: return visited
    stack = [root]
    while stack:
        node = stack.pop()
        visited += 1
        if node.left: stack.append(node.left)
        if node.right: stack.append(node.right)
    return visited

Solution 2:[2]

You can mirror the recursive algorithm using a stack:

numNodes = 0
nodeStack = [(root,0)] # (Node,0 means explore left 1 means explore right)
while nodeStack:
    nextNode, leftOrRight = nodeStack.pop()
    if not nextNode: #nextNode is empty
        continue
    if leftOrRight == 0:
        numNodes += 1
        nodeStack.append((nextNode,1))
        nodeStack.append((nextNode.leftChild,0))
    else:
        nodeStack.append((nextNode.rightChild,0))
print(numNodes)

Some things to notice: This is still a Depth-first search! That is, we still fully explore a subtree before starting to explore the other. What this means to you is that the amount of additional memory required is proportional to the height of the tree and not the width of the tree. For a balanced tree the width of the tree is 2^h where h is the height of the tree. For a totally unbalanced tree the height of the tree is the number of nodes in the tree, whereas the width is one! so it all depends on what you need :)

Now It is worth mentioning that you can make a potential optimization by checking if one of the subtrees is empty! We can change the body of if leftOrRight == 0: to:

numNodes += 1
if nextNode.rightChild: #Has a right child to explore
    nodeStack.append((nextNode,1))
nodeStack.append((nextNode.leftChild,0))

Which potentially cuts down on memory usage :)

Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source
Solution 1 EnGassa
Solution 2

Related Questions