【算法思想·搜索二叉树】构造篇

本文参考labuladong算法笔记[二叉搜索树心法（构造篇） | labuladong 的算法笔记]

1、概述

之前写了两篇手把手刷 BST 算法题的文章，第一篇 讲了中序遍历对 BST 的重要意义，第二篇 写了 BST 的基本操作。

本文就来写手把手刷 BST 系列的第三篇，循序渐进地讲两道题，如何计算所有有效 BST。

力扣 96.「不同的二叉搜索树」

给你输入一个正整数 n，请你计算，存储 {1,2,3...,n} 这些值共有多少种不同的 BST 结构。

函数签名如下：

def numTrees(n: int) -> int:

【思路】

比如说输入 n = 3，算法返回 5，因为共有如下 5 种不同的 BST 结构存储 {1,2,3}：

这就是一个正宗的穷举问题，那么什么方式能够正确地穷举有效 BST 的数量呢？

我们前文说过，不要小看「穷举」，这是一件看起来简单但是比较有技术含量的事情，问题的关键就是不能数漏，也不能数多，你咋整？

之前 手把手刷二叉树第一期 说过，二叉树算法的关键就在于明确根节点需要做什么，其实 BST 作为一种特殊的二叉树，核心思路也是一样的。

举个例子，比如给算法输入 n = 5，也就是说用 {1,2,3,4,5} 这些数字去构造 BST。

首先，这棵 BST 的根节点总共有几种情况？

显然有 5 种情况对吧，因为每个数字都可以作为根节点。

比如说我们固定 3 作为根节点，这个前提下能有几种不同的 BST 呢？

根据 BST 的特性，根节点的左子树都比根节点的值小，右子树的值都比根节点的值大。

所以如果固定 3 作为根节点，左子树节点就是 {1,2} 的组合，右子树就是 {4,5} 的组合。

左子树的组合数和右子树的组合数乘积就是 3 作为根节点时的 BST 个数。

我们这是说了 3 为根节点这一种特殊情况，其实其他的节点也是一样的。

那你可能会问，我们可以一眼看出 {1,2} 和 {4,5} 有几种组合，但是怎么让算法进行计算呢？

其实很简单，只需要递归就行了，我们可以写这样一个函数：

# 定义：闭区间 [lo, hi] 的数字能组成 count(lo, hi) 种 BST
def count(lo: int, hi: int) -> int:

【python】

注意 base case，显然当 lo > hi 闭区间 [lo, hi] 肯定是个空区间，也就对应着空节点 null，虽然是空节点，但是也是一种情况，所以要返回 1 而不能返回 0。

这样，题目的要求已经实现了，但是时间复杂度非常高，肯定存在重叠子问题。

前文动态规划相关的问题多次讲过消除重叠子问题的方法，无非就是加一个备忘录：

class Solution:
    # 备忘录
    memo = []

    def numTrees(self, n: int) -> int:
        # 备忘录的值初始化为 0
        self.memo = [[0 for _ in range(n+1)] for _ in range(n+1)]
        return self.count(1, n)

    def count(self, lo: int, hi: int) -> int:
        if lo > hi: 
            return 1
        # 查备忘录
        if self.memo[lo][hi] != 0:
            return self.memo[lo][hi]
        
        res = 0
        for mid in range(lo, hi+1):
            left = self.count(lo, mid - 1)
            right = self.count(mid + 1, hi)
            res += left * right
        # 将结果存入备忘录
        self.memo[lo][hi] = res
        
        return res

那么，如果给一个进阶题目，不止让你计算有几个不同的 BST，而是要你构建出所有有效的 BST，如何实现这个算法呢？

95. 不同的二叉搜索树 II | 力扣  | LeetCode  |

List<TreeNode> generateTrees(int n);

明白了上道题构造有效 BST 的方法，这道题的思路也是一样的：

1、穷举 root 节点的所有可能。

2、递归构造出左右子树的所有有效 BST。

3、给 root 节点穷举所有左右子树的组合。

我们可以直接看代码：

【python】

class Solution:
    def generateTrees(self, n: int) -> List[TreeNode]:
        if n == 0:
            return []
        # 构造闭区间 [1, n] 组成的 BST 
        return self.build(1, n)

    # 构造闭区间 [lo, hi] 组成的 BST 
    def build(self, lo: int, hi: int) -> List[TreeNode]:
        res = []
        # base case
        if lo > hi:
            # 这里需要装一个 null 元素，这样才能让下面的两个内层 for 循环都能进入，正确地创建出叶子节点
            # 举例来说吧，什么时候会进到这个 if 语句？当你创建叶子节点的时候，对吧。
            # 那么如果你这里不加 null，直接返回空列表，那么下面的内层两个 for 循环都无法进入        
            # 你的那个叶子节点就没有创建出来，看到了吗？所以这里要加一个 null，确保下面能把叶子节点做出来
            res.append(None)
            return res

        # 1、穷举 root 节点的所有可能。
        for i in range(lo, hi + 1):
            # 2、递归构造出左右子树的所有有效 BST。
            leftTree = self.build(lo, i - 1)
            rightTree = self.build(i + 1, hi)
            # 3、给 root 节点穷举所有左右子树的组合。
            for left in leftTree:
                for right in rightTree:
                    # i 作为根节点 root 的值
                    root = TreeNode(i)
                    root.left = left
                    root.right = right
                    res.append(root)

        return res