Quick Select

  • 找 K-th element
  • 時間複雜度: 平均 O(N)

目標

1
2
3
4
5
6
     pivot
S S S O O L L L L L

S: 小於 pivot 的數
O: 等於 pivot 的數
L: 大於 pivot 的數

過程

1
2
3
4
5
6
S S S O O X X X X L L L L L
      i   t     j

i: 指向 S 的後一個數
j: 指向 L 的前一個數
t: 指向未知的數

分堆後

1
2
3
4
5
6
S S S O O O L L L L L
a     i   j         b

if(b-j >= k) => find k-th largest in [L]
else if(b-i+1 >= k) => k-th largest is pivot
else => find "k-(b-i+1) th" largest in [S]

程式

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
int quickSelect(vector<int>& nums, int a, int b, int k){
    int pivot = nums[(a+b)/2];

    int i = a, j = b;
    int t = a;

    while(t <= j){
        if(nums[t] < pivot){
            swap(nums[i], nums[t]);
            i++;
            t++;
        }else if(nums[t] > pivot){
            swap(nums[t], nums[j]);
            j--;
        }else
            t++;
    }

    if(b-j >= k)    
        return quickSelect(nums, j+1, b, k);
    else if(b-i+1 >= k)
        return pivot;
    else
        return quickSelect(nums, a, i-1, k-(b-i+1));
}

補充

call API

1
2
3
4
5
6
7
8
9
10
11
int findKthLargest(vector<int>& nums, int k) {
    int n = nums.size();
    // nth_element 可用於求第 n 小的數 => 由小到大排
    // 第 k 大的數 => 第 n-k 小
    // nth_element 會將第n-k小的數，放在nums.begin()+n-k
    // 保證小於等於 nums.begin()+n-k 的數都在前面
    // 保證大於等於 nums.begin()+n-k 的數都在後面
    nth_element(nums.begin(), nums.begin()+n-k, nums.end());

    return *(nums.begin()+n-k);
}