LeetCode215 数组中的第K个最大元素(含堆解法实现)

发表于 | 分类于 |
字数统计: 853 | 阅读时长 ≈ 4

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

题面

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

解法

经典老番
《茴有几种写法》级别的题目

这里主要复习了一下最大堆最小堆的实现，以这道题作为背景
着重体会关于下沉时的操作 细节见代码

后面还有快排partition以及优先队列解法

代码

堆解法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
class Solution {
public:
    vector<int> heap;
    // 实现一个最小堆
    int top() {
        return heap[1];
    }

    void swim(int k) {
        while(k > 1 && heap[k / 2] > heap[k]) {
            swap(heap[k / 2], heap[k]);
            k /= 2;
        }
    }
    void push(int num) {
        heap.push_back(num);
        swim(heap.size() - 1);
    }
    void sink(int k) {
        while(k * 2 < heap.size()) {
            // 如果当前结点有子结点
            int pos = k * 2;
            if (pos + 1 < heap.size() && heap[pos + 1] < heap[pos]) {
                pos++; // 移到右子结点
            }
            if (heap[k] <= heap[pos]) { // 如果两个子结点都比父结点要大
                break; // sink结束
            }
            swap(heap[k], heap[pos]); // 交换子结点和父结点
            k = pos;
        }
    }
    void pop() {
        heap[1] = heap.back();
        heap.pop_back();
        sink(1);
    }

    int findKthLargest(vector<int>& nums, int k) {
        heap.push_back(-1);
        for (int i: nums) {
            push(i);
            if (heap.size() > k + 1) {
                pop();
            }
        }
        return top();

    }
};

集大成者

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
class Solution {
public:
    void arrangeK(vector<int>& nums, int l, int r, int k) {
        if (l > r) {
            return;
        }
        int key = nums[l], a = l, b = r;
        while (a < b) {
            while (a < b && nums[b] >= key) {
                b--;
            } 
            while(a < b && nums[a] <= key) {
                a++;
            }
            swap(nums[a], nums[b]);
        }
        swap(nums[l], nums[a]);
        // 最大的k个元素已置于右端
        if (r - a + 1 == k) {
            return;
        }
        else if (r - a + 1 < k) {
            arrangeK(nums, l, a - 1, k - (r - a + 1));
        }
        else {
            arrangeK(nums, a, r, k);
        }
    }
    int findKdc(vector<int>& nums, int k) {
        arrangeK(nums, 0, nums.size() - 1, k);
        int res = nums.back(), len = nums.size();
        for(int i = 0; i < k; i++) {
            res = min(res, nums[len - 1 - i]);
        }
        return res;
    }
    int quickSel(vector<int>& nums, int l, int r) {
        int key = nums[l];
        int a = l, b = r;
        while(a < b) {
            while(a < b && nums[b] >= key) {
                b--;
            }
            while(a < b && nums[a] <= key) {
                a++;
            }
            swap(nums[a], nums[b]);
        }
        swap(nums[a], nums[l]);
        return a;
        
    }

    int solve(vector<int>& nums, int k) {
        int l = 0, r = nums.size() - 1, target = nums.size() - k;
        // 找第k大的数也就是要有len - k个比其小的数
        while (l < r) {
            int mid = quickSel(nums, l, r);
            if (mid == target) {
                return nums[mid];
            }
            else if (mid < target) {
                l = mid + 1;
            }
            else {
                r = mid - 1;
            }
        }
        return nums[l];
    }
    int findWithPQ(vector<int>& nums, int k) {
        priority_queue<int, vector<int>, greater<int> > pq;
        for (int i: nums) {
            pq.push(i);
            if (pq.size() > k) {
                pq.pop();
            }
        }
        return pq.top();
    }
    int findKthLargest(vector<int>& nums, int k) {
        return findWithPQ(nums, k);
    }
};

partition精简版

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class Solution {
public:
    int partition(vector<int>& a, int l, int r) {
        swap(a[l], a[r]);
        int small = l;
        for (int i = l; i < r; i++) {
            if (a[i] < a[r]) {
                swap(a[i], a[small]);
                small++;
            }
        }
        swap(a[small], a[r]);
        return small;
    }
    void findk(vector<int>& nums, int l, int r, int k) {
        int index = partition(nums, l, r);
        while (index != k) {
            if (index < k) {
                l = index + 1;
            }
            else { 
                r = index - 1;
            }
            int pivot = rand() % (r - l + 1) + l;
            swap(nums[l], nums[pivot]);
            index = partition(nums, l, r);
        }
    }
    int findKthLargest(vector<int>& nums, int k) {
        int n = nums.size();
        findk(nums, 0, n - 1, n - k);
        return nums[n - k];
    }
};