[LeetCode] 215. Kth Largest Element in an Array（数组里的第 k 大元素）

• Difficulty: Medium

• Related Topics: Divide and Conqer, Heap

• Link: https://leetcode.com/problems/kth-largest-element-in-an-array/

Description

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.
在一个无序数组中找出第 K 大元素。注意这里的第 K 大指的是排序后的第 K 大元素，重复元素被包括在内。

Examples

Example 1

Input: [3,2,1,5,6,4] and k = 2
Output: 5

Example 2

Input: [3,2,3,1,2,4,5,5,6] and k = 4
Output: 4

Note

You may assue k is always valid, 1 ≤ k ≤ array's length.

Solution

其实这题最粗暴的方法就是调用一个 sort 然后取出第 k 个，不过这么做显然就没什么意思了。这里提供一个使用堆（优先队列）的方法：维护一个容量为 k 的堆，将数组中的元素依次入堆，若堆容量超了就弹出一个元素，最后堆顶元素即为所求，代码如下：

import java.util.*

class Solution {
    fun findKthLargest(nums: IntArray, k: Int): Int {
        val q = PriorityQueue<Int>()

        for (num in nums) {
            q.offer(num)
            if (q.size > k) {
                q.poll()
            }
        }

        return q.peek()
    }
}

不过我印象中，在算法课上有学过一个快速选择的算法，该算法来源于快排，基本思想是分治，但快排需要应付左右两侧，快速选择每次二分之后，只需要对其中一侧继续处理即可。代码我忘了这么写，以下代码来自于 discussion：

class Solution {
    fun findKthLargest(nums: IntArray, k: Int): Int {
        var start = 0
        var end = nums.lastIndex
        val index = nums.size - k

        while (start < end) {
            val pivot = partition(nums, start, end)
            when {
                pivot < index -> start = pivot + 1
                pivot > index -> end = pivot - 1
                else -> return nums[pivot]
            }
        }
        return nums[start]
    }

    private fun partition(nums: IntArray, start: Int, endInclusive: Int): Int {
        val pivot = start
        var p = start
        var q = endInclusive
        while (p <= q) {
            while (p <= q && nums[p] <= nums[pivot]) {
                p++
            }
            while (p <= q && nums[q] > nums[pivot]) {
                q--
            }
            if (p > q) {
                break
            }
            nums.swap(p, q)
        }
        nums.swap(q, pivot)
        return q
    }

    private fun IntArray.swap(i: Int, j: Int) {
        val t = this[i]
        this[i] = this[j]
        this[j] = t
    }
}