[LintCode] 第k大元素
基于快速排序:
1 class Solution { 2 public: 3 /* 4 * param k : description of k 5 * param nums : description of array and index 0 ~ n-1 6 * return: description of return 7 */ 8 int kthLargestElement(int k, vector<int> nums) { 9 // write your code here 10 int left = 0, right = nums.size() - 1; 11 while (true) { 12 int pos = partition(nums, left, right); 13 if (pos == k - 1) return nums[pos]; 14 if (pos < k - 1) left = pos + 1; 15 if (pos > k - 1) right = pos - 1; 16 } 17 } 18 private: 19 int partition(vector<int>& nums, int left, int right) { 20 int pivot = nums[left]; 21 int l = left + 1, r = right; 22 while (l <= r) { 23 if (nums[l] < pivot && nums[r] > pivot) 24 swap(nums[l++], nums[r--]); 25 if (nums[l] >= pivot) l++; 26 if (nums[r] <= pivot) r--; 27 } 28 swap(nums[left], nums[r]); 29 return r; 30 } 31 };
基于最大堆:
1 class Solution { 2 public: 3 /* 4 * param k : description of k 5 * param nums : description of array and index 0 ~ n-1 6 * return: description of return 7 */ 8 inline int left(int idx) { 9 return (idx << 1) + 1; 10 } 11 inline int right(int idx) { 12 return (idx << 1) + 2; 13 } 14 void max_heapify(vector<int>& nums, int idx) { 15 int largest = idx; 16 int l = left(idx), r = right(idx); 17 if (l < heap_size && nums[l] > nums[largest]) 18 largest = l; 19 if (r < heap_size && nums[r] > nums[largest]) 20 largest = r; 21 if (largest != idx) { 22 swap(nums[idx], nums[largest]); 23 max_heapify(nums, largest); 24 } 25 } 26 void build_max_heap(vector<int>& nums) { 27 heap_size = nums.size(); 28 for (int i = (heap_size >> 1) - 1; i >= 0; i--) 29 max_heapify(nums, i); 30 } 31 int kthLargestElement(int k, vector<int> nums) { 32 // write your code here 33 build_max_heap(nums); 34 for (int i = 0; i < k; i++) { 35 swap(nums[0], nums[heap_size - 1]); 36 heap_size--; 37 max_heapify(nums, 0); 38 } 39 return nums[heap_size]; 40 } 41 private: 42 int heap_size; 43 };