LeetCode Notes_#703_数据流的第K大元素

LeetCode

Contents

题目
思路分析
方法1：二叉搜索树
方法2：小顶堆
解答
解答1：二叉搜索树
复杂度分析
解答2：小顶堆
复杂度分析

题目

设计一个找到数据流中第K大元素的类（class）。注意是排序后的第K大元素，不是第K个不同的元素。
你的 KthLargest 类需要一个同时接收整数 k 和整数数组nums 的构造器，它包含数据流中的初始元素。每次调用 KthLargest.add，返回当前数据流中第K大的元素。

示例:

int k = 3;
int[] arr = [4,5,8,2];
KthLargest kthLargest = new KthLargest(3, arr);
kthLargest.add(3);   // returns 4
kthLargest.add(5);   // returns 5
kthLargest.add(10);  // returns 5
kthLargest.add(9);   // returns 8
kthLargest.add(4);   // returns 8

说明:
你可以假设 nums 的长度≥ k-1 且k ≥ 1。

思路分析

这题有两个思路

用二叉搜索树保存数据。
用大小为k的小顶堆保存数据，堆顶的数据就是要找的第k大数字。

方法1：二叉搜索树

首先需要定义一下BST的节点类。注意需要增加一个count来记录子树节点个数。

//内部类，定义TreeNode
    class TreeNode{
        int val;
        int count = 1;//记录当前节点子树的节点个数（包括自己）
        TreeNode left;
        TreeNode right;
        TreeNode() {};
        TreeNode(int val) {this.val = val;}
    }

然后实现BST的两个辅助方法，可参考 《算法第四版》3.2节。

private TreeNode insertIntoBST(TreeNode root, int val)向BST插入一个节点
private int selectKthNum(TreeNode root, int K),搜索BST里第k大的数字

最后就很容易实现了。

构造器KthLargest(int k, int[] nums)
- 将nums中的数字全部插入BST。
add()方法
- 将值为val的节点插入BST。
- 用selectKthNum方法搜索出第k大的数字。

方法2：小顶堆

思路比较简单，详见代码。

解答

解答1：二叉搜索树

class KthLargest {
    //内部类，定义TreeNode
    class TreeNode{
        int val;
        int count = 1;//记录当前节点子树的节点个数（包括自己）
        TreeNode left;
        TreeNode right;
        TreeNode() {};
        TreeNode(int val) {this.val = val;}
    }

    int K;
    TreeNode root = null;

    //向BST当中插入一个节点
    private TreeNode insertIntoBST(TreeNode root, int val){
        TreeNode node = root;
        while(node != null){
            //题目没有说加入的数字都是不相同的，那么将相等的数字放在左边
            if(val <= node.val){
                node.count += 1;
                if(node.left == null){
                    node.left = new TreeNode(val);
                    return root;
                }
                else node = node.left;
            }
            else if(val > node.val){
                node.count += 1;
                if(node.right == null){
                    node.right = new TreeNode(val);
                    return root;
                }
                else node = node.right;
            }
        }
        return new TreeNode(val);
    }

    //参考《算法第四版》3.2.3.4
    //寻找二叉搜索树当中第K大的节点,这个节点的右子树size应该是K-1
    private int selectKthNum(TreeNode root, int K){
        //返回-1代表遍历到叶子节点还是没找到符合要求的
        if(root == null) return -1;
        int rightTreeSize = 0;
        if(root.right != null) rightTreeSize = root.right.count;
        //如果root.right的节点个数恰好是K-1,返回这个节点的值
        if(rightTreeSize == K -1) return root.val;
        //如果当前节点右子树节点个数大于K-1,说明要找的节点在root.right当中
        else if(rightTreeSize > K - 1) return selectKthNum(root.right, K);
        //如果当前节点右子树节点个数小于K-1,说明要找的节点在root.left当中
        //需要减去rightTreeSize + 1,因为已经有这么多比较大的数字了
        else return selectKthNum(root.left, K - rightTreeSize - 1);
    }

    public KthLargest(int k, int[] nums) {
        //赋值给全局变量，以便于add()方法使用
        K = k;
        //将nums当中所有数字插入BST
        for(int i = 0;i <= nums.length - 1;i++){
            root = insertIntoBST(root, nums[i]);
        }
    }

    public int add(int val) {
        //1.插入值为val的节点
        root = insertIntoBST(root, val);
        //2.寻找第k大数字，然后返回
        return selectKthNum(root, K);
    }
}

复杂度分析

时间复杂度：
插入操作和搜索第k大数字，复杂度都是O(h),h = logn。

构造器KthLargest:O(nlogn)
add方法:O(logn)

空间复杂度：O(n),使用n个TreeNode。

解答2：小顶堆

class KthLargest {
    //优先队列实现小顶堆
    private PriorityQueue<Integer> queue;
    private int limit;

    public KthLargest(int k, int[] nums) {
        limit = k;
        queue = new PriorityQueue<>(k);
        //将所有元素加入小顶堆
        for (int num : nums) {
            add(num);
        }
    }

    public int add(int val) {
        //小顶堆还没填满，直接加入val
        if (queue.size() < limit) {
            queue.add(val);
        //小顶堆已经满了，且val比堆内的最小值大，先弹出最小值，再加入val
        }else if (val > queue.peek()) {
            queue.poll();
            queue.add(val);
        }
        //返回堆顶数据即可
        return queue.peek();
    }
}

复杂度分析

PriorityQueue:插入O(logn),取出O(1)
时间复杂度：

构造器KthLargest:O(nlogn)
add方法:O(logn)

空间复杂度：O(k)。