diff --git a/problems/0110.平衡二叉树.md b/problems/0110.平衡二叉树.md index d98ff8a9..1b997643 100644 --- a/problems/0110.平衡二叉树.md +++ b/problems/0110.平衡二叉树.md @@ -531,40 +531,26 @@ class Solution: 迭代法: ```python class Solution: - def isBalanced(self, root: TreeNode) -> bool: - st = [] + def isBalanced(self, root: Optional[TreeNode]) -> bool: if not root: return True - st.append(root) - while st: - node = st.pop() #中 - if abs(self.getDepth(node.left) - self.getDepth(node.right)) > 1: - return False - if node.right: - st.append(node.right) #右(空节点不入栈) - if node.left: - st.append(node.left) #左(空节点不入栈) - return True - - def getDepth(self, cur): - st = [] - if cur: - st.append(cur) - depth = 0 - result = 0 - while st: - node = st.pop() + + height_map = {} + stack = [root] + while stack: + node = stack.pop() if node: - st.append(node) #中 - st.append(None) - depth += 1 - if node.right: st.append(node.right) #右 - if node.left: st.append(node.left) #左 + stack.append(node) + stack.append(None) + if node.left: stack.append(node.left) + if node.right: stack.append(node.right) else: - node = st.pop() - depth -= 1 - result = max(result, depth) - return result + real_node = stack.pop() + left, right = height_map.get(real_node.left, 0), height_map.get(real_node.right, 0) + if abs(left - right) > 1: + return False + height_map[real_node] = 1 + max(left, right) + return True ```