Balanced Binary Tree LeetCode solution in Python

1
# Definition for a binary tree node.
2
class TreeNode:
3
    def __init__(self, val=0, left=None, right=None):
4
        self.val = val
5
        self.left = left
6
        self.right = right
7

8

9
class Solution:
10
    def isBalanced(self, root: TreeNode) -> bool:
11

12
        # Initialize the result to True
13
        res = True
14

15
        # DFS through the tree
16
        def dfs(node, i):
17
            nonlocal res
18

19
            # If there isn't a node, return previous depth
20
            if not node:
21
                return i - 1
22

23
            # Check depths of the left and right subtrees
24
            left, right = dfs(node.left, i + 1), dfs(node.right, i + 1)
25

26
            # If they are more than 1 difference, save False to the result
27
            if abs(right - left) > 1:
28
                res = False
29

30
            # Return the max depth of both subtrees
31
            return max(left, right)
32

33
        dfs(root, 0)
34

35
        return res