LeetCode 二叉搜索树系列 ～ 数据流中的第 K 大元素

携手创作，共同成长！这是我参与「掘金日新计划 · 8 月更文挑战」的第10天，点击查看活动详情

二叉搜索树的定义

二叉查找树（Binary Search Tree），（又：二叉搜索树，二叉排序树）它或者是一棵空树，或者是具有下列性质的二叉树： 若它的左子树不空，则左子树上所有结点的值均小于它的根结点的值； 若它的右子树不空，则右子树上所有结点的值均大于它的根结点的值； 它的左、右子树也分别为二叉排序树。二叉搜索树作为一种经典的数据结构，它既有链表的快速插入与删除操作的特点，又有数组快速查找的优势；所以应用十分广泛，例如在文件系统和数据库系统一般会采用这种数据结构进行高效率的排序与检索操作。

二叉搜索树主要支持三个操作

搜索

关于二叉搜索树的搜索，对于每个节点：

• 如果目标值等于节点的值，则返回节点；
• 如果目标值小于节点的值，则继续在左子树中搜索；
• 如果目标值大于节点的值，则继续在右子树中搜索。

插入

插入到操作与搜索类似，对于每个节点：

• 根据节点值与目标节点值的关系，搜索左子树或右子树；
• 重复上一步直到到达外部节点；
• 根据节点的值与目标节点的值的关系，将新节点添加为其左侧或右侧的子节点。

删除

删除的操作比起前两个操作会稍微复杂些，需要注意的点如下：

• 如果目标节点没有子节点，我们可以直接移除该目标节点。
• 如果目标节只有一个子节点，我们可以用其子节点作为替换。
• 如果目标节点有两个子节点，我们需要用其中序后继节点或者前驱节点来替换，再删除该目标节点。

题目

设计一个找到数据流中第 k 大元素的类（class）。注意是排序后的第 k 大元素，不是第 k 个不同的元素。

请实现 KthLargest 类：

• KthLargest(int k, int[] nums) 使用整数 k 和整数流 nums 初始化对象。
• int add(int val) 将 val 插入数据流 nums 后，返回当前数据流中第 k 大的元素。

示例

输入：
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
输出：
[null, 4, 5, 5, 8, 8]

解释：
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3);   // return 4
kthLargest.add(5);   // return 5
kthLargest.add(10);  // return 5
kthLargest.add(9);   // return 8
kthLargest.add(4);   // return 8

提示

• 1 <= k <= 10^4
• 0 <= nums.length <= 10^4
• -10^4 <= nums[i] <= 10^4
• -10^4 <= val <= 10^4
• 最多调用 add 方法 10^4 次
• 题目数据保证，在查找第 k 大元素时，数组中至少有 k 个元素

数据流中的第 K 大元素代码实现

class KthLargest {
    private int k;
    private BST bst = new BST();

    public KthLargest(int k, int[] nums) {
        this.k = k;
        // 将数组元素构建成二叉树
        for (int n : nums) {
            bst.add(n);
        }
    }
    
    public int add(int val) {
        bst.add(val);
        return bst.search(k).val;
    }
    
    private class BST {
        private TreeNode root;

        public void add(int val) {
            root = add(root, val);
        }

        // 二叉树构建
        private TreeNode add(TreeNode node, int val) {
            if (node == null) {
                return new TreeNode(val);
            }
            if (node.val > val) {
                node.left = add(node.left, val);
            } else if (node.val < val) {
                node.right = add(node.right, val);
            }
            // 元素重复 不添加进树但是count++
            node.count++;
            return node;
        }
        
        public TreeNode search(int k) {
            return search(root, k);
        }

        // 目标值搜索
        private TreeNode search(TreeNode node, int k) {
            if (node == null) {
                return node;
            }
            int leftNodeCount = node.left != null ? node.left.count : 0;
            int rightNodeCount = node.right != null ? node.right.count : 0;
            int currNodeCount = node.count - leftNodeCount - rightNodeCount;
            if (k > currNodeCount + rightNodeCount ) {
                // k > 当前结点数加右子树的结点数，则搜索左子树
                return search(node.left, k - currNodeCount - rightNodeCount);
            } else if (k <= rightNodeCount) {
                // k <= 右子树的结点数，则搜索右子树
                return search(node.right, k);
            } else {
                // k == 当前结点数加右子树的结点数，则找到第k大的结点
                return node;
            }
        }

        // 自定义二叉树
        private class TreeNode {
            // 节点值
            private int val;
            // 左节点
            private TreeNode left;
            // 右节点
            private TreeNode right;
            // 统计数量
            private int count;

            TreeNode(int val) { 
                this.val = val;
                this.count = 1;
            }
        }
    }
}