大根堆

FlyingLight發表於2024-08-07

/*

  • @lc app=leetcode.cn id=215 lang=java
  • [215] 陣列中的第K個最大元素
    */

// @lc code=start
class Solution {
private int[] smallHeap;
int len = 0;

void swap(int idx1, int idx2) {
    int tmp = smallHeap[idx1];
    smallHeap[idx1] = smallHeap[idx2];
    smallHeap[idx2] = tmp;
}

boolean cmp(int idx1, int idx2) { return smallHeap[idx1] > smallHeap[idx2]; }

void heapify(int idx) {
    if (idx < 0 || idx >= len) return;
    int l = (idx << 1) | 1, r = (idx << 1) + 2;
    if (l >= len) return;
    if (r == len) {
        if (cmp(idx, l))
            swap(idx, l);
        return;
    }
    if (cmp(l, idx) && cmp(r, idx)) return;
    int swapIdx = (cmp(idx, l) && cmp(r, l)) ? l : r;
    swap(swapIdx, idx);
    heapify(swapIdx);
}

int poll() {
    int res = smallHeap[0];
    smallHeap[0] = smallHeap[--len];
    smallHeap[len] = 0;
    heapify(0);
    return res;
}

int peek() { return smallHeap[0]; }

void add(int x) {
    int idx = len;
    smallHeap[len++] = x;
    while(idx > 0 && cmp(idx >> 1, idx)) {
        swap(idx >> 1, idx);
        idx >>= 1;
    }
}

public int findKthLargest(int[] nums, int k) {
    smallHeap = new int[nums.length];
    for (int i = 0; i < k; ++i) add(nums[i]);
    for (int i = k; i < nums.length; ++i) {
        if (nums[i] > peek()) {
            add(nums[i]);
            poll();
        }
    }

    return peek();
}

}
// @lc code=end

相關文章