215. Kth Largest Element in an Array

, ,

Priority Queue

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Solution {
    public int findKthLargest(int[] nums, int k) {
        Queue<Integer> minHeap = new PriorityQueue<>(); // 小顶堆，小于堆顶的都被过滤了，剩下的就是 topK
        for (int i = 0; i < k; i++) {
            minHeap.offer(nums[i]);
        }

        for (int i = k; i < nums.length; i++) {
            if (nums[i] > minHeap.peek()) { // if 判断不要亦可，具有该判断时性能更优，即只将比堆顶数值更大的数值入堆
                minHeap.offer(nums[i]);
                minHeap.poll();
            }
        }

        return minHeap.peek();
    }
}

QuickSort

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
class Solution {
    public int findKthLargest(int[] nums, int k) {
        int targetIndex = k - 1;

        int left = 0, right = nums.length - 1;
        while (true) {
            int pivotIndex = partition(nums, left, right);
            if (pivotIndex < targetIndex) {
                left = pivotIndex + 1;
            } else if (pivotIndex > targetIndex) {
                right = pivotIndex - 1;
            } else {
                return nums[pivotIndex];
            }
        }
    }

    private int partition(int[] nums, int left, int right) {
        // 随机选择一个元素交换至最左侧，性能优化使用，亦可不交换
        int randomIndex = left + ThreadLocalRandom.current().nextInt(right - left + 1);
        swap(nums, left, randomIndex);

        int pivotValue = nums[left];
        int endIndex = left; // 指向左边界（含）元素
        for (int i = left + 1; i <= right; i++) {
            if (nums[i] >= pivotValue) {
                swap(nums, i, ++endIndex);
            }
        }
        // nums[left + 1, ... endIndex] >= nums[left]
        swap(nums, left, endIndex);
        return endIndex;
    }

    private void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}

QuickSort + Two Pointers

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
class Solution {
    public int findKthLargest(int[] nums, int k) {
        int targetIndex = k - 1;

        int left = 0, right = nums.length - 1;
        while (true) {
            int pivotIndex = partition(nums, left, right);
            if (pivotIndex < targetIndex) {
                left = pivotIndex + 1;
            } else if (pivotIndex > targetIndex) {
                right = pivotIndex - 1;
            } else {
                return nums[targetIndex];
            }
        }
    }

    private int partition(int[] nums, int left, int right) {
        int randomIndex = left + ThreadLocalRandom.current().nextInt(right - left + 1);
        int pivotValue = nums[randomIndex];
        swap(nums, left, randomIndex);

        int i = left, j = right;
        while (i < j) {
            while (i < j && nums[j] <= pivotValue) {
                j--;
            }
            while (i < j && nums[i] >= pivotValue) {
                i++;
            }

            swap(nums, i, j);
        }

        swap(nums, left, i); // 注意此处是 i
        return i; // 注意此处是 i
    }

    private void swap(int[] nums, int i, int j) {
        int tmp = nums[i];
        nums[i] = nums[j];
        nums[j] = tmp;
    }
}

注意此解法在数组含有大量重复元素时性能会大幅下降，时间复杂度降低为 O(n^2)，所以有了下面的解法。

Quick Select

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
class Solution {
    public int findKthLargest(int[] nums, int k) {
        List<Integer> numList = new ArrayList<>(nums.length);
        for (int num : nums) {
            numList.add(num);
        }

        return findKthLargest(numList, k);
    }

    private int findKthLargest(List<Integer> numList, int k) {
        List<Integer> largeList = new ArrayList<>();
        List<Integer> equalList = new ArrayList<>();
        List<Integer> smallList = new ArrayList<>();

        int pivotValue = numList.get(ThreadLocalRandom.current().nextInt(numList.size()));

        for (int num : numList) {
            if (num > pivotValue) {
                largeList.add(num);
            } else if (num == pivotValue) {
                equalList.add(num);
            } else {
                smallList.add(num);
            }
        }

        if (largeList.size() >= k) {
            return findKthLargest(largeList, k);
        } else if (largeList.size() + equalList.size() >= k) {
            return pivotValue;
        } else {
            return findKthLargest(smallList, k - largeList.size() - equalList.size());
        } 
    }
}

该解法保证了重复元素过多时的处理效率，时间复杂度为 O(N)，空间复杂度为 O(logN)。

References

215. Kth Largest Element in an Array
剑指 Offer II 076. 数组中的第 k 大的数字