Search code examples
pythondepth-first-search

Max() function on node ; python DFS


in the classic example of DFS below:

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

def tree_max_depth(root: Node) -> int:
    def dfs(root):
        # null node adds no depth
        if not root:
            return 0
        # num nodes in longest path of current subtree = max num nodes of its two subtrees + 1 current node
        return max(dfs(root.left), dfs(root.right)) + 1
    print(dfs(root.left))
    return dfs(root) - 1 if root else 0

# this function builds a tree from input; you don't have to modify it
# learn more about how trees are encoded in https://algo.monster/problems/serializing_tree
def build_tree(nodes, f):
    val = next(nodes)
    if val == 'x': return None
    left = build_tree(nodes, f)
    right = build_tree(nodes, f)
    return Node(f(val), left, right)

if __name__ == '__main__':
    root = build_tree(iter(input().split()), int)
    res = tree_max_depth(root)
    print(res)

How is the max function calculating the height in the line below?

return max(dfs(root.left), dfs(root.right)) + 1

Solution

  • If you understand recursion, then you can understand what the block of code does.

    Let's say initial node is A. In the begining, max(dfs(root.left), dfs(root.right)) + 1 is like saying max(dfs(nodeB), dfs(nodeC)) + 1.

    • First, it calculates dfs(nodeB). This has root.right/left, so it doesn't return 0.

    • It continues to max(dfs(nodeD), dfs(nodeE)) + 1.

    • It goes to NodeD (root.left) which has None as root.left/right, so it returns `max(dfs(None), dfs(None))+1)=max(0,0)+1=0+1=1)

    • Then, it continues to dfs(nodeE) (the root.right). which have root.left/right. So it goes to dfs(nodeF). dfs(nodeF) is root has NOT root.left/right so it returns 1 (like nodeD).

    So now, we have for the B node, the code max(1, 2) where '1' is from root.left (nodeD) and '2' is from nodes E and F. Then chooses the max, which is 2, and then returns 2+1 to the node A.

    Node C, has height 1. So, max(3,1) is 1. So, max height is 3.

    An explanation of the root

    Each time dfs(root) is called, then a new root is set. For example, if we call dfs(nodeB) (assume nodeB is a node), then the new root is nodeB.

    In the block of code:

    if not root:
        return 0
    

    the if-statement, in reality checks if the root at the moment, is None. If we substitute the code above with a root being None, it is:

    if not None:
        return 0
    

    and not None as we know results to true.

    Let's take NodeD. When it's root, it does NOT have root.right & root.left (they are None). So it is calling max(dfs(None), dfs(None)+1 which equals to max(0, 0)+1 because dfs(None) returns 0.

    enter image description here