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

什么是平衡二叉搜索树？

平衡二叉搜索树（Balanced Binary Search Tree，简称BBST）是一种特殊的二叉搜索树，它具有以下几个属性：

• 每个节点的左子树和右子树的高度绝对值不超过1，也就是说它们的深度相差不超过1；
• 所有叶子节点到根节点的路径长度相同；
• 它是一棵二叉搜索树，也就是左子树包含的所有节点的值小于节点的值，右子树包含的所有节点的值大于节点的值。

因为平衡二叉搜索树的这些属性，它在插入、删除、查询等操作中的时间复杂度是O(log n)的，而不需要在最坏情况下变成链表时的O(n)时间复杂度。

平衡二叉搜索树常用的实现方式

常见的平衡二叉搜索树实现方式包括：

• AVL树：由G.M. Adelson-Velsky和E.M. Landis于1962年提出，是第一种平衡二叉树；
• 红黑树：由R. Bayer和E.M. McCreight于1972年提出，是应用最广泛的平衡二叉树之一；
• Treap：由C.R. Aragon和R. Seidel于1996年提出，是一种随机化数据结构，它在平均情况下的性能优于AVL树和红黑树。

问题 13：找到二叉搜索树中第K小的元素

问题描述：给定一棵二叉搜索树，请找出其中第k小的元素。

解法1：中序遍历

二叉搜索树的中序遍历是有序的，所以可以先进行中序遍历，然后找到第k小的元素即可。时间复杂度为O(n)，空间复杂度为O(n)，其中n是二叉搜索树中的节点数。

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        def inorder(node):
            if not node:
                return []
            return inorder(node.left) + [node.val] + inorder(node.right)
    
        return inorder(root)[k-1]

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        List<Integer> inorder = inorderTraversal(root);
        return inorder.get(k-1);
    }
    
    private List<Integer> inorderTraversal(TreeNode node) {
        List<Integer> res = new ArrayList<>();
        if (node == null) {
            return res;
        }
        res.addAll(inorderTraversal(node.left));
        res.add(node.val);
        res.addAll(inorderTraversal(node.right));
        return res;
    }
}

解法2：优化的中序遍历

由于题目要求的是第k小的元素，我们可以在遍历过程中提前终止，不必继续遍历整棵树。时间复杂度可以降低到O(k)，空间复杂度仍然是O(n)。

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        def inorder(node):
            if not node or len(res) == k:
                return
            inorder(node.left)
            if len(res) == k:
                return
            res.append(node.val)
            inorder(node.right)
        
        res = []
        inorder(root)
        return res[-1]

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        List<Integer> res = new ArrayList<>();
        inorderTraversal(root, res, k);
        return res.get(k-1);
    }
    
    private void inorderTraversal(TreeNode node, List<Integer> res, int k) {
        if (node == null || res.size() == k) {
            return;
        }
        inorderTraversal(node.left, res, k);
        if (res.size() == k) {
            return;
        }
        res.add(node.val);
        inorderTraversal(node.right, res, k);
    }
}

解法3：使用堆

我们可以使用一个小根堆来维护当前遍历到的k个最小的元素。遍历完所有节点后，堆顶元素就是第k小的元素。时间复杂度为O(n log k)，空间复杂度为O(k)。

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        heap = []
        def inorder(node):
            if not node:
                return
            if len(heap) < k:
                heapq.heappush(heap, -node.val)
            else:
                if node.val < -heap[0]:
                    heapq.heappushpop(heap, -node.val)
            inorder(node.left)
            inorder(node.right)
    
        inorder(root)
        return -heap[0]

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>();
        inorderTraversal(root, heap, k);
        return heap.peek();
    }
    
    private void inorderTraversal(TreeNode node, PriorityQueue<Integer> heap, int k) {
        if (node == null) {
            return;
        }
        if (heap.size() < k) {
            heap.offer(node.val);
        } else {
            if (node.val > heap.peek()) {
                heap.poll();
                heap.offer(node.val);
            }
        }
        inorderTraversal(node.left, heap, k);
        inorderTraversal(node.right, heap, k);
    }
}