题目描述

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

示例：
输入: [3,2,1,5,6,4], k = 2
输出: 5

题目分析

优先队列

遍历数组，放入小顶优先队列，超出 k 个时取出1个顶部元素。
最终顶部元素即为第K个最大元素。
时间复杂度 O(n * logk)

哈希计数

遍历一边数组，使用哈希表记录每个数字出现的次数，同时找出最大值和最小值。
从最大值开始对次数求和，sum >= k 时退出。
时间复杂度 O(n)

快速选择

快排的过程中判断 k 落在哪个区间，然后递归.
每层平均减少一半元素，处理次数为 n + n/2 + n/4 … + n/x <= 2n
时间复杂度 O(n)

Java

public int findKthLargest(int[] nums, int k) {
    PriorityQueue<Integer> queue = new PriorityQueue<>();
    for (int x : nums) {
        queue.offer(x);
        if (queue.size() > k) {
            queue.poll();
        }
    }
    return queue.poll();
}

Kotlin

// 优先队列
fun findKthLargest(nums: IntArray, k: Int): Int {
    val queue = PriorityQueue<Int>()
    for (n in nums) {
        queue.offer(n)
        if (queue.size > k) {
            queue.poll()
        }
    }
    return queue.poll()
}

// 哈希表计数
fun findKthLargest2(nums: IntArray, k: Int): Int {
    var min = nums.first()
    var max = nums.first()
    val count = HashMap<Int, Int>()
    for (n in nums) {
        count[n] = (count[n] ?: 0) + 1
        min = min(min, n)
        max = max(max, n)
    }
    var sum = 0
    for (n in max downTo min) {
        sum += count[n] ?: 0
        if (sum >= k) {
            return n
        }
    }
    return 0
}


// 快速选择
fun findKthLargest(nums: IntArray, k: Int, start: Int = 0, end: Int = nums.size - 1): Int {
    val target = nums[end]
    var index = start
    for (i in start..end) {
        if (nums[i] > target) {
            swap(nums, i, index++)
        }
    }
    swap(nums, end, index)
    val count = index - start // 左侧数量

    return if (count < k - 1) {
        findKthLargest(nums, k - count - 1, index + 1, end)
    } else if (count > k - 1) {
        findKthLargest(nums, k, start, index - 1)
    } else {
        nums[index]
    }
}

private fun swap(nums: IntArray, i: Int, j: Int) {
    val temp = nums[i]
    nums[i] = nums[j]
    nums[j] = temp
}