数据结构 | 平衡二叉搜索树 | 问题2

什么是平衡二叉搜索树？

平衡二叉搜索树（Balanced Binary Search Tree），简称为平衡树（Balanced Tree）是一种特殊的二叉搜索树，它通过调整树中各个节点的左右子树的高度，使得整个树的高度尽量小，从而提高树的查找、插入和删除等操作的效率。

常见的平衡二叉搜索树有AVL树、红黑树、B树/B+树等。

什么是问题2？

问题2是指在平衡二叉搜索树中，找到排名第k小的元素。

具体来说，如果平衡二叉搜索树中有n个元素，那么排名第k小的元素就是第k个被访问到的元素，对于二叉搜索树的中序遍历，就是第k个被遍历到的元素。

如何解决问题2？

方法1：暴力解法

对平衡二叉搜索树进行中序遍历，记录当前访问到第几个元素，如果第k个元素被访问到，返回该元素的值即可。

时间复杂度：O(n)，n为二叉搜索树中元素的个数，最坏情况下需要访问所有的元素。

代码片段：

class Solution:
    def __init__(self):
        self.count = 0
        self.res = None
    
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        self.inorderTraversal(root, k)
        return self.res
    
    def inorderTraversal(self, root, k):
        if not root:
            return
        self.inorderTraversal(root.left, k)
        self.count += 1
        if self.count == k:
            self.res = root.val
            return
        self.inorderTraversal(root.right, k)

方法2：二分法

可以利用二分法来查找第k小的元素。

对于二叉搜索树的每个节点，可以计算出其左子树中元素的个数left和右子树中元素的个数right。对于当前节点A，若k <= left，则说明要查找的元素应该在A的左子树中，否则，如果k > left + 1，则说明要查找的元素应该在A的右子树中，此时更新k值为k - left - 1，继续在右子树中查找。

时间复杂度：O(logn)，n为二叉搜索树中元素的个数，由于每个节点都需要遍历一遍，最多需要遍历二叉搜索树的高度，即logn。

代码片段：

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        left_count = self.getCount(root.left)
        if k <= left_count:
            return self.kthSmallest(root.left, k)
        elif k == left_count + 1:
            return root.val
        else:
            return self.kthSmallest(root.right, k-left_count-1)
        
    def getCount(self, root):
        if not root:
            return 0
        return self.getCount(root.left) + self.getCount(root.right) + 1

总结

问题2是平衡二叉搜索树中比较常见的问题，解决方法主要有暴力解法和二分法，其中二分法的时间复杂度比较优秀，可以达到O(logn)的时间复杂度。