215. 数组中的第K个最大元素

 1 #include <algorithm>
 2 #include <iostream>
 3 #include <vector>
 4 
 5 class Solution {
 6  public:
 7   int findKthLargest(std::vector<int>& nums, int k) {
 8     if (k < 1) {
 9       return -1;
10     }
11     int begin = 0, end = nums.size() - 1;
12     int mid = -1;
13     while (mid + 1 != k) {
14       mid = partition(nums, begin, end);
15       if (mid + 1 > k) {
16         end = mid - 1;
17       } else if (mid + 1 < k) {
18         begin = mid + 1;
19       }
20     }
21 
22     return nums[mid];
23   }
24 
25  private:
26   int partition(std::vector<int>& nums, int begin, int end) {
27     if (begin > end || end > nums.size() || nums.empty()) {
28       return -1;
29     }
30 
31     int key = nums[begin];
32     int i = begin, j = end;
33     while (i < j) {
34       while (i < j && nums[j] <= key) --j;
35       while (i < j && nums[i] >= key) ++i;
36       std::swap(nums[i], nums[j]);
37     }
38 
39     std::swap(nums[i], nums[begin]);
40     return i;
41   }
42 };
43 
44 int main() {
45   std::vector<int> nums = {3, 2, 3, 1, 2, 4, 5, 5, 6};
46   Solution so;
47   std::cout << so.findKthLargest(nums, 4) << std::endl;
48   return 0;
49 }

  1 #include <algorithm>
  2 #include <iostream>
  3 #include <vector>
  4 
  5 class LittleHeap {
  6  public:
  7   void push(int n) {
  8     heap_.push_back(n);
  9     adjustUp();
 10   }
 11 
 12   void pop() {
 13     if (heap_.empty()) {
 14       return;
 15     }
 16 
 17     std::swap(heap_.front(), heap_.back());
 18     heap_.pop_back();
 19     adjustDown(0);
 20   }
 21 
 22   int top() {
 23     if (heap_.empty()) {
 24       return -1;
 25     }
 26 
 27     return heap_.front();
 28   }
 29 
 30   // debug
 31   void debugString() {
 32     std::for_each(heap_.begin(), heap_.end(),
 33                   [](int n) { std::cout << n << ‘ ‘; });
 34     std::cout << std::endl;
 35   }
 36 
 37   uint32_t size() { return heap_.size(); }
 38 
 39  private:
 40   // i = 0,1,2,3,4...
 41   // the child of a[i] is a[2*i + 1], a[2*i + 2]
 42   void adjustDown(int root) {
 43     if (heap_.size() < 2) {
 44       return;
 45     }
 46 
 47     int child = 2 * root + 1;
 48     while (child < heap_.size()) {
 49       if (child + 1 < heap_.size() && heap_[child + 1] < heap_[child]) {
 50         ++child;
 51       }
 52 
 53       if (heap_[root] > heap_[child]) {
 54         std::swap(heap_[root], heap_[child]);
 55         root = child;
 56         child = 2 * root + 1;
 57       } else {
 58         break;
 59       }
 60     }
 61   }
 62 
 63   // the parent of a[i] is a[(i-1) / 2]
 64   void adjustUp() {
 65     if (heap_.size() < 2) {
 66       return;
 67     }
 68 
 69     int child = heap_.size() - 1;
 70     int parent = (child - 1) / 2;
 71     while (parent >= 0) {
 72       if (heap_[parent] > heap_[child]) {
 73         std::swap(heap_[parent], heap_[child]);
 74         child = parent;
 75         parent = (child - 1) / 2;
 76       } else {
 77         break;
 78       }
 79     }
 80   }
 81 
 82  private:
 83   std::vector<int> heap_;
 84 };
 85 
 86 class Solution {
 87  public:
 88   int findKthLargest(std::vector<int>& nums, int k) {
 89     LittleHeap heap;
 90     for (const auto& i : nums) {
 91       heap.push(i);
 92       if (heap.size() > k) {
 93         heap.pop();
 94       }
 95     }
 96     return heap.top();
 97   }
 98 };
 99 
100 int main() {
101   std::vector<int> nums = {3, 2, 3, 1, 2, 4, 5, 5, 6};
102   Solution so;
103   std::cout << so.findKthLargest(nums, 4) << std::endl;
104   return 0;
105 }