题目

给定一个二叉树,判断它是否是高度平衡的二叉树。

本题中,一棵高度平衡二叉树定义为:

一个二叉树每个节点 的左右两个子树的高度差的绝对值不超过1。

示例 1:

给定二叉树 [3, 9, 20, null, null, 15, 7]

  1. 3
  2. / \
  3. 9 20
  4. / \
  5. 15 7

返回 true

示例 2:

给定二叉树 [1, 2, 2, 3, 3, null, null, 4, 4]

       1
      / \
     2   2
    / \
   3   3
  / \
 4   4

返回 false

方案一

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def isBalanced(self, root: TreeNode) -> bool:
        if not root:
            return True
        elif root.left and root.right:
            if abs(self._getHeight(root.left) - self._getHeight(root.right)) > 1:
                return False
            else:
                return self.isBalanced(root.left) and self.isBalanced(root.right)
        elif root.left:
            return self._getHeight(root.left) == 1
        elif root.right:
            return self._getHeight(root.right) == 1
        else:
            return True

    def _getHeight(self, root, height=1):
        '''获取树的高度'''
        if not root:
            return 0
        elif root.left and root.right:
            return max(self._getHeight(root.left, height+1), self._getHeight(root.right, height+1))
        elif root.left:
            return self._getHeight(root.left, height+1)
        elif root.right:
            return self._getHeight(root.right, height+1)
        else:
            return height
  • _getHeight 的时间复杂度为 平衡二叉树 - 图1,这就导致了 isBalanced 时间复杂度为 平衡二叉树 - 图2
  • _getHeight 的写法又导致了无法使用 @functools.lru_cache

因为每次调用 _getHeight 参数不同,导致不好进行缓存。

修改后如下:

import functools

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def isBalanced(self, root: TreeNode) -> bool:
        # 查看缓存使用情况
        # print(self._getHeight.cache_info())

        if not root:
            return True
        elif root.left and root.right:
            if abs(self._getHeight(root.left) - self._getHeight(root.right)) > 1:
                return False
            else:
                return self.isBalanced(root.left) and self.isBalanced(root.right)
        elif root.left:
            return self._getHeight(root.left) == 1
        elif root.right:
            return self._getHeight(root.right) == 1
        else:
            return True        

    @functools.lru_cache
    def _getHeight(self, root):
        '''
        @return height
        '''
        if not root:
            return 0
        elif root.left and root.right:
            return max(self._getHeight(root.left), self._getHeight(root.right)) + 1
        elif root.left:
            return self._getHeight(root.left) + 1
        elif root.right:
            return self._getHeight(root.right) + 1
        else:
            return 1

参考 leetcode 简化代码后写法

class Solution:
    def isBalanced(self, root: TreeNode) -> bool:
        if not root:
            return True

        l = self._getHeight(root.left)
        r = self._getHeight(root.right)

        return abs(l - r) <= 1 and self.isBalanced(root.left) and self.isBalanced(root.right)

    @functools.lru_cache
    def _getHeight(self, root):
        '''
        @return height
        '''
        if not root:
            return 0

        l = self._getHeight(root.left)
        r = self._getHeight(root.right)

        return max(l, r) + 1

方案二(leetcode)

class Solution:
    def isBalanced(self, root: TreeNode) -> bool:
        return self.depth(root) != -1

    def depth(self, root):
        if not root: return 0
        left = self.depth(root.left)
        if left == -1: return -1
        right = self.depth(root.right)
        if right == -1: return -1
        return max(left, right) + 1 if abs(left - right) < 2 else -1
  • 注意最后一个 return,当不满足条件是可以直接返回,可以降低时间复杂度;
  • 相比我上述的 方案一 而言,对此题效率更高,但是通用性不如 方案一

原文

https://leetcode-cn.com/explore/learn/card/introduction-to-data-structure-binary-search-tree/67/appendix-height-balanced-bst/188/