[LintCode]k番目の大きな要素

11097 ワード

クイック・ソート・ベース:
 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 };