48 lines
1.8 KiB
Python
48 lines
1.8 KiB
Python
|
# Time: O(N)
|
||
|
|
# Space: O(N)
|
||
|
|
|
||
|
|
# Definition for a binary tree node.
|
||
|
|
# class TreeNode:
|
||
|
|
# def __init__(self, val=0, left=None, right=None):
|
||
|
|
# self.val = val
|
||
|
|
# self.left = left
|
||
|
|
# self.right = right
|
||
|
|
from collections import namedtuple
|
||
|
|
|
||
|
|
class Solution:
|
||
|
|
def isBalanced(self, root: Optional[TreeNode]) -> bool:
|
||
|
|
Result = namedtuple('Result', 'is_balanced height')
|
||
|
|
|
||
|
|
def dfs(node):
|
||
|
|
"""
|
||
|
|
Instead of doing top-down, we'll do bottom-up recursion
|
||
|
|
via DFS to solve subproblems and bubble back up to the root
|
||
|
|
"""
|
||
|
|
|
||
|
|
# This happens when we reach leaf node, in which case, we assume
|
||
|
|
# things are balanced and return 0 height
|
||
|
|
if node is None: return Result(True, 0)
|
||
|
|
|
||
|
|
# DFS recursion
|
||
|
|
right = dfs(node.right)
|
||
|
|
left = dfs(node.left)
|
||
|
|
|
||
|
|
# For current `node`, things are only going to be balanced if
|
||
|
|
# both left and right subtrees are balanced. Otherwise, we can
|
||
|
|
# return False right away.
|
||
|
|
has_balanced_subtrees = right.is_balanced and left.is_balanced
|
||
|
|
|
||
|
|
# Besides having left and right subtrees themselves *individually*
|
||
|
|
# being balanced, we need to next check height difference <= 1.
|
||
|
|
if has_balanced_subtrees and abs(right.height - left.height) <= 1:
|
||
|
|
# Height of tree formed by current `node` would be the max
|
||
|
|
# height of its left/right subtree + 1 (itself)
|
||
|
|
return Result(True, 1 + max(left.height, right.height))
|
||
|
|
|
||
|
|
# If it reaches here, that means either height diff > 1 or left/right
|
||
|
|
# subtrees are already imbalanced.
|
||
|
|
return Result(False, 0)
|
||
|
|
|
||
|
|
|
||
|
|
return dfs(root).is_balanced
|