Python: pass “mutable integer” in recursion

by

I’m working on some code to find the max depth/height in a tree. I came up with a recursive solution, and I find myself needing to pass an integer “by reference”. Currently, I’m using a mutable container to achieve this, but is there a more Pythonic way?

class Node:
"""
Node in a tree, with children stored in a list.
"""
   def __init__(self, keyParam, childrenParam = None):
    self.key = keyParam #some id of the node; not strictly required, but often needed.

       if childrenParam is None:
           self.children = []
       else:
           self.children = childrenParam

def dfsWithDepth(node, currentDepth, currentMaxContainer):
    currentDepth += 1
    if currentDepth > currentMaxContainer[0]:
        currentMaxContainer[0] = currentDepth

    for child in node.children:
        if child is not None:
            dfsWithDepth(child, currentDepth, currentMaxContainer)

# %% Construct tree
rootNode = Node('n1')
rootNode.children = [Node('n2'), Node('n3')]
tempNode = rootNode.children[0]  # n2
tempNode.children = [Node('n3'), Node('n4')]
tempNode = tempNode.children[1]  # n4
tempNode.children = [None, Node('n5')]

currentMaxContainer = [0]
dfsWithDepth(rootNode, 0, currentMaxContainer)

print(currentMaxContainer)

Answer

Optimization and restructuring

Node class constructor (__init__ method)

  • the arguments keyParam, childrenParam introduce unneeded verbosity and better named as just key and children
  • the whole if ... else ... condition for ensuring children list is simply replaced with self.children = children or [] expression

The initial dfsWithDepth function should be renamed to say find_max_depth to follow PEP8 style guide naming conventions.

The former parameters currentDepth, currentMaxContainer become unneeded.
Within the crucial recursive function it’s worth to capture “empty” node and node with no children beforehand. That’s achieved by the following condition:

if node is None or not node.children:
    return 1

The remaining recursive search relies on the crucial find_max_depth and builtin max functions calls and accumulates the final max depth through traversing all subsequent sub-trees (child nodes).

The final optimized version:

class Node:
    """
    Node in a tree, with children stored in a list.
    """

    def __init__(self, key, children=None):
        self.key = key  # some id of the node; not strictly required, but often needed.
        self.children = children or []


def find_max_depth(node):
    if node is None or not node.children:
        return 1
    return 1 + max(map(find_max_depth, node.children))

Testing (I’ve added additional 5th level for demonstration):

# Construct tree
rootNode = Node('n1')
rootNode.children = [Node('n2'), Node('n3')]
tempNode = rootNode.children[0]  # n2
tempNode.children = [Node('n3'), Node('n4')]
tempNode = tempNode.children[1]  # n4
tempNode.children = [None, Node('n5')]
tempNode = tempNode.children[1]  # n5
tempNode.children = [Node('name7')]

print(find_max_depth(rootNode))   # 5

Attribution
Source : Link , Question Author : rishai , Answer Author : RomanPerekhrest

Leave a Comment