设计一个找到数据流中第 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 <= 104

  • 0 <= nums.length <= 104

  • -104 <= nums[i] <= 104

  • -104 <= val <= 104

  • 最多调用 add 方法 104

  • 题目数据保证,在查找第 k 大元素时,数组中至少有 k 个元素

思路

  • 优先级队列构建堆排序

  • 手动实现堆排序

  • 复习堆排序知识

代码

  • 优先级队列构建堆排序

class KthLargest {
    private PriorityQueue<Integer> pq;
    private int k;

    public KthLargest(int k, int[] nums) {
        this.pq = new PriorityQueue<>();
        this.k = k;
        for (int num : nums) {
            doAdd(num);
        }
    }

    private void doAdd(int num) {
        if (pq.size() < k) {
            pq.offer(num);
        } else if (pq.peek() < num) {
            pq.offer(num);
            pq.poll();
        }
    }
    
    public int add(int val) {
        doAdd(val);
        return pq.peek();
    }
}
  • 手动实现堆排序

class KthLargest {
    private int[] heap;  // 堆数组
    private int size;    // 当前堆中的元素个数
    private int k;      // 第k大元素

    public KthLargest(int k, int[] nums) {
        this.k = k;
        this.heap = new int[k];  // 初始化堆数组
        this.size = 0;           // 初始堆为空

        // 逐个添加初始元素
        for (int num : nums) {
            add(num);
        }
    }

    public int add(int val) {
        // 如果堆未满,直接插入并调整
        if (size < k) {
            heap[size] = val;    // 将新元素放在末尾
            swim(size);          // 上浮调整
            size++;              // 堆大小增加
        }
        // 堆已满且新元素大于堆顶(最小元素)
        else if (val > heap[0]) {
            heap[0] = val;       // 替换堆顶
            sink(0);             // 下沉调整
        }
        //printHeap();
        // 返回堆顶元素(即第k大元素)
        return heap[0];
    }

    // 上浮操作:将索引为i的元素向上调整
    private void swim(int i) {
        while (i > 0) {
            int parent = (i - 1) / 2;  // 父节点索引
            // 如果当前节点小于父节点,则交换
            if (heap[i] < heap[parent]) {
                swap(i, parent);
                i = parent;
            } else {
                break;  // 调整完成
            }
        }
    }

    // 下沉操作:将索引为i的元素向下调整
    private void sink(int i) {
        while (i < size) {
            int left = 2 * i + 1;   // 左孩子索引
            int right = 2 * i + 2;  // 右孩子索引
            int smallest = i;       // 记录最小元素的索引

            // 与左孩子比较
            if (left < size && heap[left] < heap[smallest]) {
                smallest = left;
            }
            // 与右孩子比较
            if (right < size && heap[right] < heap[smallest]) {
                smallest = right;
            }

            // 如果当前节点是最小的,调整结束
            if (smallest == i) {
                break;
            }

            // 否则交换并继续下沉
            swap(i, smallest);
            i = smallest;
        }
    }

    // 交换堆中两个元素
    private void swap(int i, int j) {
        int temp = heap[i];
        heap[i] = heap[j];
        heap[j] = temp;
    }
}

复习堆排序知识

堆排序中的‌向上调整‌和‌向下调整‌是两种不同的建堆方式,主要区别在于调整方向和起始节点:

调整方向

  • 向上调整‌:从第二个节点开始,逐步向上比较并调整父节点,直到根节点

  • 向下调整‌:从最后一个非叶子节点开始,逐步向下比较并调整子节点

起始节点

  • 向上调整通常从第二个节点开始(因为第一个节点是根节点,无需向上调整)

  • 向下调整从最后一个非叶子节点开始(通常为数组末尾的父节点)

效率与适用场景

  • 向上调整‌:时间复杂度与向下调整相同,本质都是插入建堆,但未额外开辟空间

  • 向下调整‌:更适合处理动态插入数据的情况,通过从底部开始调整可避免多次遍历父节点

注意!!!向下调整算法有一个前提,左右子树必须是一个堆,才能进行调整

两种方法均可实现堆排序,但具体选择需结合数据插入顺序和性能需求。