[算法基础] 堆相关-数组中第K个最大元素

题目

在未排序的数组中找到第 k 个最大的元素。请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

示例 1:

输入: [3,2,1,5,6,4] 和 k = 2
输出: 5

示例 2:

输入: [3,2,3,1,2,4,5,5,6] 和 k = 4
输出: 4

说明:

你可以假设 k 总是有效的，且 1 ≤ k ≤ 数组的长度。

解题思路

我们可以用小顶堆来解决这个问题。截取数组前k个元素构建一个小顶堆，循环剩余部分，遇到比堆顶大的元素，就替换掉堆顶元素。最终这个小顶堆中的元素就是整个数组中最大的那k个元素，而堆顶元素即为这k个元素中最小的那一个。

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var findKthLargest = function(nums, k) {
	// 截取前k个数，构建一个小顶堆
	var heap = [,], i = 0
	while(i < k){
		heap.push(nums[i++]
	}
	bulidHeap(heap, k)
	// 从 k 位开始遍历数组
	for(var i = k;i < nums.length; i++){
		if(heap[1] < nums[i]){
		// 如果小顶堆的顶部元素小于遍历到的数，则替换并重新堆化
			heap[1] = nums[i]
			heapify(heap, k ,1)
		}
		// 返回堆顶部元素
		return heap[1]
	}
};

// 原地建堆
var buildHeap = function(arr, k){
    if(k === 1) return
    // 从最后一个非叶子节点开始，自上而下式堆化
    for(let i = Math.floor(k/2); i>=1 ; i--) {
        heapify(arr, k, i)
    }
}

// 堆化
var heapify = function(arr, k, i){
    // 自上而下式堆化
    while(true) {
        let minIndex = i
        if(2*i <= k && arr[2*i] < arr[i]) {
            minIndex = 2*i
        }
        if(2*i+1 <= k && arr[2*i+1] < arr[minIndex]) {
            minIndex = 2*i+1
        }
        if(minIndex !== i) {
            swap(arr, i, minIndex)
            // 循环条件
            i = minIndex
        } else {
            break
        }
    }
}

// 交换
let swap = (arr, i , j) => {
    let temp = arr[i]
    arr[i] = arr[j]
    arr[j] = temp
}