【LeetCode】703.数据流中的第K大元素（二叉搜索树解法,java实现）

题目

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

分析

这种方法的时空复杂度不是最优解，但是可以练习熟悉二叉搜索树的功能。

算法步骤：

（1）：为创建二叉搜索树做准备，设置一个二叉树节点的内部类用来存储

• 该节点的值val;
• 该节点所在的子树的节点数数量cnt（包括自己，cnt>=1）;
• 该节点的左子树右子树对象left,right;
• 内部类的构造体，用于在建立新的对象时初始化该节点的数据；

（2）：将整数数组数据存储为二叉搜索树，使用全局变量存储二叉树根节点和 K（第k大元素）的值；
（3）：每添加一个数时，先用二叉搜索树插入方法将这个数值加入树中，然后判断第k大元素的数值：

利用迭代的方法查找：

• 判断节点右子树节点数量是否大于k,如果大于，则开始在右子树查找；
• 如果小于，则k值减去（右子树节点树+1），开始在左子树查找；直到k == 1时结束，返回此时节点的值。

代码实现：

public class KthLargest {
    private int key;
    private TreeNode root;

    //内部类
    private class TreeNode{
        private int val;
        private int cnt;
        private TreeNode left;
        private TreeNode right;
        TreeNode() {}
        TreeNode(int x){
            val = x;
            cnt = 1;
        }
    }

    private TreeNode insertNode(TreeNode root, int val){
        if(root == null){
            return new TreeNode(val);
        }

        if(val <= root.val){
            root.cnt++;
            root.left = insertNode(root.left, val);
        }else{
            root.cnt++;
            root.right = insertNode(root.right, val);
        }
        return root;
    }

    public KthLargest(int k, int[] nums) {
        key = k;
        for(int i : nums){
            root = insertNode(root, i);
        }
    }

    public int add(int val) {
        root = insertNode(root, val);
        int k = key;
        TreeNode node = root;

        while(k >= 1) {
            if (node.right != null) {
                if ((node.right.cnt + 1) > k) {
                    node = node.right;
                } else if ((node.right.cnt + 1) < k) {
                    k = k - node.right.cnt - 1;
                    node = node.left;
                } else {
                    return node.val;
                }
            } else {
                if (k == 1) {
                    return node.val;
                } else {
                    k--;
                    node = node.left;
                }
            }
        }
        return node.val;
        }
}