【LeetCode 二叉树专项】二叉搜索树中第 K 小的元素（230）

文章目录

1. 题目

1.1 示例

1.2 说明

1.3 限制

1.4 进阶

2. 解法一（递归中序遍历）

2.1 分析

2.2 实现

2.3 复杂度

3. 解法二（迭代中序遍历）

3.1 分析

3.2 实现

3.3 复杂度

4. 解法三（记录子树的节点数量）

4.1 分析

4.2 解答

4.3 复杂度

5. 解法四（平衡二叉搜索树）

1. 题目

给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 $1$ 开始计数）。

1.1 示例

示例 $1$ ：
- 输入： root = [3, 1, 4, null, 2] ， k = 1
- 输出： $1$

在这里插入图片描述

示例 $2$ ：
- 输入： root = [5, 3, 6, 2, 4, null, null, 1], k = 3
- 输出： $3$

在这里插入图片描述

1.2 说明

来源： 力扣（LeetCode）

链接： https://leetcode-cn.com/problems/kth-smallest-element-in-a-bst

1.3 限制

树中的节点数为 $n$ ；
$\le k \le n \le 10^4$ ；
$\le Node.val \le 10^4$ 。

1.4 进阶

如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化算法？

2. 解法一（递归中序遍历）

2.1 分析

实际上，由于题目指出这里给定的是一棵二叉搜索树，而由【数据结构Python描述】二叉搜索树简介与使用二叉搜索树实现有序映射可知，二叉搜索树有一个特殊的性质，即其中序遍历的结果是一个有序的递增序列，利用此性质可以方便地给出如下递归解法：

2.2 实现

from typing import Optional, List


class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def _kth_smallest(self, root: Optional[TreeNode], tree: List[int]):
        if not root:
            return
        self._kth_smallest(root.left, tree)
        tree.append(root.val)
        self._kth_smallest(root.right, tree)

    def kth_smallest(self, root: Optional[TreeNode], k: int) -> Optional[int]:
        tree = []
        self._kth_smallest(root, tree)
        print(tree)
        return tree[k - 1]


def main():
    node6 = TreeNode(1)
    node5 = TreeNode(4)
    node4 = TreeNode(2, left=node6)
    node3 = TreeNode(6)
    node2 = TreeNode(3, left=node4, right=node5)
    node1 = TreeNode(5, left=node2, right=node3)
    root = node1
    sln = Solution()
    print(sln.kth_smallest(root, 3))  # 3


if __name__ == '__main__':
    main()

实际上，这里也可以无需将整个二叉搜索树的所有节点遍历后再根据 k 的值得到结果，可以通过维护一个实例属性 self._counter 来标示中序遍历到了第几小的元素：

from typing import Optional, List


class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def __init__(self):
        self.result = 0
        self._counter = 0

    def _efficient_kth_smallest(self, root: Optional[TreeNode], k: int):
        if not root:
            return
        self._efficient_kth_smallest(root.left, k)
        self._counter += 1
        if self._counter == k:
            self.result = root.val
            return
        self._efficient_kth_smallest(root.right, k)

    def efficient_kth_smallest(self, root: Optional[TreeNode], k: int) -> Optional[int]:
        self._efficient_kth_smallest(root, k)
        return self.result


def main():
    node6 = TreeNode(1)
    node5 = TreeNode(4)
    node4 = TreeNode(2, left=node6)
    node3 = TreeNode(6)
    node2 = TreeNode(3, left=node4, right=node5)
    node1 = TreeNode(5, left=node2, right=node3)
    root = node1
    sln = Solution()
    print(sln.efficient_kth_smallest(root, 3))  # 3


if __name__ == '__main__':
    main()

2.3 复杂度

时间复杂度： $O (n)$ ；

空间复杂度： $O (n)$ 。

3. 解法二（迭代中序遍历）

3.1 分析

由于二叉树除了递归方式的中序遍历，还有迭代方式的中序遍历，因此可通过稍加修改迭代的形式的二叉树中序遍历得到如下解法：

3.2 实现

from typing import Optional, List


class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def iterative_kth_smallest(self, root: Optional[TreeNode], k: int) -> Optional[int]:
        if not root:
            return
        stack = []
        cursor = root
        result = 0
        while True:
            if cursor:
                stack.append(cursor)
                cursor = cursor.left
            elif stack:
                node = stack.pop()
                k -= 1
                if k == 0:
                    result = node.val
                    break
                cursor = node.right
            else:
                break
        return result


def main():
    node6 = TreeNode(1)
    node5 = TreeNode(4)
    node4 = TreeNode(2, left=node6)
    node3 = TreeNode(6)
    node2 = TreeNode(3, left=node4, right=node5)
    node1 = TreeNode(5, left=node2, right=node3)
    root = node1
    sln = Solution()
    print(sln.iterative_kth_smallest(root, 3))  # 3


if __name__ == '__main__':
    main()

3.3 复杂度

时间复杂度： $O (n)$ ；

空间复杂度： $O (n)$ 。

4. 解法三（记录子树的节点数量）

4.1 分析

如果你需要频繁地查找第 k 小的值，你将如何优化算法？

在上述解法中，我们之所以需要中序遍历前 k 个元素，是因为我们不知道子树的结点数量，不得不通过遍历子树的方式来获知。

因此，我们可以记录下以每个结点为根结点的子树的结点数，并在查找第 k 小的值时，使用如下方法搜索：

令 node 等于根结点，开始搜索。对当前结点 node 进行如下操作：

如果 node 的左子树的结点数 left 小于 k − 1 ，则第 k 小的元素一定在 node 的右子树中，令 node 等于其的右子结点， k 等于 k − left − 1 ，并继续搜索；
如果 node 的左子树的结点数left 等于 k − 1，则第 k 小的元素即为 node ，结束搜索并返回 node.val 即可；
如果 node 的左子树的结点数left 大于 k − 1 ，则第 k 小的元素一定在 node 的左子树中，令 node 等于其左子结点，并继续搜索。

在实现时，我们既可以将以每个结点为根结点的子树的结点数存储在结点中，也可以将其记录在哈希表中。

作者： LeetCode-Solution

4.2 解答

下面实现是基于将以每个节点为根节点的子树的节点数存储在节点对象中给出的实现：

from typing import Optional


class SizeableTreeNode:
    def __init__(self, val=0, left: 'SizeableTreeNode' = None, right: 'SizeableTreeNode' = None, size: 'int' = 0):
        self.val = val
        self.left = left
        self.right = right
        self.size = size


class Solution:
    def __init__(self):
        self.result = float('-inf')

    def _kth_smallest(self, root: Optional[SizeableTreeNode], k: int):
        if not root:
            return
        if root.left:
            size = root.left.size
        else:
            size = 0
        if size < k - 1:
            self._kth_smallest(root.right, k - size - 1)
        elif size == k - 1:
            self.result = root.val
            return
        else:
            self._kth_smallest(root.left, k)

    def kth_smallest(self, root: Optional[SizeableTreeNode], k: int):
        self._kth_smallest(root, k)
        if self.result == float('-inf'):
            return
        else:
            return self.result


def test_sizeable_node():
    node6 = SizeableTreeNode(1, size=1)
    node5 = SizeableTreeNode(4, size=1)
    node4 = SizeableTreeNode(2, left=node6, size=2)
    node3 = SizeableTreeNode(6, size=1)
    node2 = SizeableTreeNode(3, left=node4, right=node5, size=4)
    node1 = SizeableTreeNode(5, left=node2, right=node3, size=6)
    root = node1
    sln = Solution()
    print(sln.kth_smallest(root, 7))  # None
    print(sln.kth_smallest(root, 6))  # 6


def main():
    test_sizeable_node()


if __name__ == '__main__':
    main()

下面实现是基于将以每个节点为根节点的子树的节点数存储在哈希表中（这里使用了 Python 的 collections 模块中专门用于统计值为非负整型数值的 Counter 类）给出的实现：

from collections import Counter
from typing import Optional


class TreeNode:
    def __init__(self, val=0, left: 'TreeNode' = None, right: 'TreeNode' = None):
        self.val = val
        self.left = left
        self.right = right


class SizeableBinarySearchTree:
    def __init__(self, root: TreeNode):
        self.root = root
        self._node_num = Counter()
        self._count_num(root)

    def _count_num(self, root: TreeNode) -> int:
        if not root:
            return 0
        self._node_num[root] = 1 + self._count_num(root.left) + self._count_num(root.right)
        return self._node_num[root]

    def _get_num(self, root: TreeNode) -> int:
        if root:
            return self._node_num[root]
        else:
            return 0

    def kth_smallest(self, k: int):
        node = self.root
        while node:
            left = self._get_num(node.left)
            if left < k - 1:
                node = node.right
                k = k - 1 - left
            elif left == k - 1:
                return node.val
            else:
                node = node.left


class Solution:
    def freq_search_kth_smallest(self, root: Optional[TreeNode], k: int):
        bst = SizeableBinarySearchTree(root)
        return bst.kth_smallest(k)


def test_sizeable_binary_search_tree():
    node6 = TreeNode(1)
    node5 = TreeNode(4)
    node4 = TreeNode(2, left=node6)
    node3 = TreeNode(6)
    node2 = TreeNode(3, left=node4, right=node5)
    node1 = TreeNode(5, left=node2, right=node3)
    root = node1
    sln = Solution()
    print(sln.freq_search_kth_smallest(root, 7))  # None
    print(sln.freq_search_kth_smallest(root, 6))  # 6


def main():
    test_sizeable_binary_search_tree()


if __name__ == '__main__':
    main()

4.3 复杂度

时间复杂度： 预处理的时间复杂度为 $O (n)$ ，其中 $n$ 是树中结点的总数；我们需要遍历树中所有结点来统计以每个结点为根结点的子树的结点数。频繁搜索的时间复杂度为 $O (h)$ ，其中 $h$ 是树的高度；当树是平衡树时，时间复杂度取得最小值 $O(\log n)$ ；当树是线性树时，时间复杂度取得最大值 $O (n)$ 。
空间复杂度： $O (n)$ ，用于存储以每个结点为根结点的子树的结点数。

5. 解法四（平衡二叉搜索树）

作者：LeetCode-Solution