分治题目：数组中的第 K 个最大元素

文章目录

• 题目
• • 标题和出处
  • 难度
  • 题目描述
  • • 要求
    • 示例
    • 数据范围
• 解法
• • 思路和算法
  • 代码
  • 复杂度分析

题目

标题和出处

标题：数组中的第 K 个最大元素

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

难度

6 级

题目描述

要求

给定整数数组 nums \texttt{nums} nums 和整数 k \texttt{k} k，返回数组中第 k \texttt{k} k 个最大的元素。

注意需要返回的是数组排序后的第 k \texttt{k} k 个最大的元素，不是第 k \texttt{k} k 个不同的元素。

要求时间复杂度是 O(n) \texttt{O(n)} O(n)。

示例

示例 1：

输入： nums = 3,2,1,5,6,4, k = 2 \texttt{nums = 3,2,1,5,6,4, k = 2} nums = 3,2,1,5,6,4, k = 2

输出： 5 \texttt{5} 5

示例 2：

输入： nums = 3,2,3,1,2,4,5,5,6, k = 4 \texttt{nums = 3,2,3,1,2,4,5,5,6, k = 4} nums = 3,2,3,1,2,4,5,5,6, k = 4

输出： 4 \texttt{4} 4

数据范围

• 1 ≤ k ≤ nums.length ≤ 10 5 \texttt{1} \le \texttt{k} \le \texttt{nums.length} \le \texttt{10}^\texttt{5} 1≤k≤nums.length≤105
• -10 4 ≤ numsi ≤ 10 4 \texttt{-10}^\texttt{4} \le \texttt{numsi} \le \texttt{10}^\texttt{4} -104≤numsi≤104

解法

思路和算法

在长度为 n n n 的数组中寻找第 k k k 个最大的元素，最简单的做法是将数组排序，然后返回第 k k k 个最大的元素，该做法的时间复杂度是 O ( n log ⁡ n ) O(n \log n) O(nlogn)。

为了将时间复杂度降低到 O ( n ) O(n) O(n)，需要使用快速选择算法。快速选择算法和快速排序算法相似，是分治的应用。由于题目要求寻找第 k k k 个最大的元素，因此考虑将数组降序排序，排序后的数组的下标 k − 1 k - 1 k−1 处元素即为第 k k k 个最大的元素。

快速选择算法的操作包括选择基准元素和分区，做法如下。

1. 选择基准元素。可以选择数组的首个元素作为基准元素，也可以在数组中随机选择一个元素作为基准元素并将基准元素与数组的首个元素交换位置，此时数组的首个元素为基准元素。

2. 分区。分区的含义是将数组分成两个子数组，基准元素左侧子数组的元素都大于等于基准元素，基准元素右侧子数组的元素都小于基准元素。具体做法如下。

   1. 用 start \textit{start} start 和 end \textit{end} end 分别表示当前数组的开始下标和结束下标，初始化两个指针 low = start + 1 \textit{low} = \textit{start} + 1 low=start+1 和 high = end \textit{high} = \textit{end} high=end。

   2. 将 low \textit{low} low 从左往右遍历，直到遇到小于基准元素的元素；将 high \textit{high} high 从右往左遍历，直到遇到大于等于基准元素的元素。此时如果 low < high \textit{low} < \textit{high} low<high，则交换 low \textit{low} low 和 high \textit{high} high 处的元素。重复该操作，直到 low ≥ high \textit{low} \ge \textit{high} low≥high 时结束该操作。

   3. 将 high \textit{high} high 从右往左遍历，直到遇到大于等于基准元素的元素。

   4. 此时 high \textit{high} high 指向基准元素应该放置下标的位置。如果 high > start \textit{high} > \textit{start} high>start，则交换 start \textit{start} start 和 high \textit{high} high 处的元素。

3. 用 pivotIndex \textit{pivotIndex} pivotIndex 表示分区之后基准元素所在下标。根据 pivotIndex \textit{pivotIndex} pivotIndex 与 k − 1 k - 1 k−1 的大小关系，执行如下操作。

   • 如果 pivotIndex = k − 1 \textit{pivotIndex} = k - 1 pivotIndex=k−1，则此时的基准元素即为第 k k k 个最大的元素，返回下标 pivotIndex \textit{pivotIndex} pivotIndex 处的元素。

   • 如果 pivotIndex > k − 1 \textit{pivotIndex} > k - 1 pivotIndex>k−1，则下标 pivotIndex \textit{pivotIndex} pivotIndex 处的元素小于等于第 k k k 个最大的元素，在下标范围 start , pivotIndex − 1 \\textit{start}, \\textit{pivotIndex} - 1 start,pivotIndex−1 中继续寻找第 k k k 个最大的元素。如果数组中存在多个基准元素，则可以跳过连续基准元素，从而降低时间复杂度。

   • 如果 pivotIndex < k − 1 \textit{pivotIndex} < k - 1 pivotIndex<k−1，则下标 pivotIndex \textit{pivotIndex} pivotIndex 处的元素大于等于第 k k k 个最大的元素，在下标范围 pivotIndex + 1 , end \\textit{pivotIndex} + 1, \\textit{end} pivotIndex+1,end 中继续寻找第 k k k 个最大的元素。

快速排序算法的平均时间复杂度是 O ( n log ⁡ n ) O(n \log n) O(nlogn)，快速选择算法在每次分区之后都可以排除不可能包含第 k k k 个最大元素的子数组，因此平均时间复杂度低于快速排序算法。

平均情况下，快速选择算法在每次分区之后都可以排除数组中的一半元素，递归调用层数是 O ( log ⁡ n ) O(\log n) O(logn)，时间复杂度是 O ( n ) O(n) O(n)，空间复杂度是 O ( log ⁡ n ) O(\log n) O(logn)。最差情况下，快速选择算法在每次分区时选择的基准元素都是数组中的最小值或最大值，递归调用层数是 O ( n ) O(n) O(n)，时间复杂度是 O ( n 2 ) O(n^2) O(n2)，空间复杂度是 O ( n ) O(n) O(n)。

使用随机选择基准元素的做法可以最大程度避免最差情况的发生，达到 O ( n ) O(n) O(n) 的时间复杂度。

代码

class Solution {
    public int findKthLargest(int[] nums, int k) {
        return quickSelect(nums, k - 1, 0, nums.length - 1);
    }

    public int quickSelect(int[] nums, int index, int start, int end) {
        int pivotIndex = partition(nums, start, end);
        if (pivotIndex == index) {
            return nums[pivotIndex];
        } else if (pivotIndex > index) {
            while (pivotIndex - 1 > index && nums[pivotIndex - 1] == nums[pivotIndex]) {
                pivotIndex--;
            }
            return quickSelect(nums, index, start, pivotIndex - 1);
        } else {
            return quickSelect(nums, index, pivotIndex + 1, end);
        }
    }

    public int partition(int[] nums, int start, int end) {
        int randomIndex = start + (int) (Math.random() * (end - start + 1));
        swap(nums, start, randomIndex);
        int pivot = nums[start];
        int low = start, high = end;
        while (low < high) {
            while (low < high && nums[low] >= pivot) {
                low++;
            }
            while (low < high && nums[high] < pivot) {
                high--;
            }
            if (low < high) {
                swap(nums, low, high);
            }
        }
        while (high > start && nums[high] < pivot) {
            high--;
        }
        if (high > start) {
            swap(nums, start, high);
        }
        return high;
    }

    public void swap(int[] nums, int index1, int index2) {
        int temp = nums[index1];
        nums[index1] = nums[index2];
        nums[index2] = temp;
    }
}

复杂度分析

• 时间复杂度：平均情况是 O ( n ) O(n) O(n)，最差情况是 O ( n 2 ) O(n^2) O(n2)，其中 n n n 是数组 nums \textit{nums} nums 的长度。快速选择算法的平均时间复杂度是 O ( n ) O(n) O(n)，最差时间复杂度是 O ( n 2 ) O(n^2) O(n2)。

• 空间复杂度：平均情况是 O ( log ⁡ n ) O(\log n) O(logn)，最差情况是 O ( n ) O(n) O(n)，其中 n n n 是数组 nums \textit{nums} nums 的长度。空间复杂度取决于递归调用层数。平均情况下，递归调用层数是 O ( log ⁡ n ) O(\log n) O(logn)，快速选择算法的空间复杂度是 O ( log ⁡ n ) O(\log n) O(logn)。最差情况下，递归调用层数是 O ( n ) O(n) O(n)，快速选择算法的空间复杂度是 O ( n ) O(n) O(n)。