215.数组中第K个最大元素

102 阅读4分钟

给定整数数组 nums 和整数 k,请返回数组中第 k 个最大的元素。

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

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

 

示例 1:

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

示例 2:

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

 

提示:

  • 1 <= k <= nums.length <= 105
  • -104 <= nums[i] <= 104

题目思路

本题是非常好的一题,可以同时应用两个经典排序算法:

  • 快排(利用 PartitionK 算法快速筛选出第 K 大元素)
    • 分区算法1:普通方法,按照快排的思路进行分区,先从后往前找第一个比 pivot 小的元素,然后从前往后找第一个比 pivot 大的元素,交换;重复这个过程,最终双指针中左侧指针 i 指向 pivot 的最终位置,而右侧指针 j 则指向 pivot 位置的前一个位置(因为 j 从右往左找的过程中,只会在最终小于 pivot 的地方停下),最终分区的位置可以直接返回 i 也可以返回 j+1;
    • 分区算法2:取中法,为了使分区的效果更好,选择左侧、中间、右侧元素中的中间元素作为 pivot。
    • 分区算法3:普通方法的变形,将 pivot 交换到末尾,维护一个比 pivot 小的指针,从左到右遍历数组元素,发现比 pivot 小的元素,交换到 start 指针指向的位置,指针继续向前移动,最终指针刚好指向 pivot 的最终位置,交换后返回指针位置即可。
    • 特殊算法:以上三种方法都无法处理中间有大量重复元素的数组排列(如[0,1,1,1,1,1,1,1,1,2,3,4,1,1,1,1,1,1,5,3,2,1,1]),表现为40/41 用例无法通过,解决方法为强制移动。
  • 堆排(维护一个大小为 K 的小根堆,当所有元素处理完之后,小根堆中剩余的就是数组中最大的 K 个元素了,此时小根堆堆顶就是第 K 大元素了)

PartitionK

func findKthLargest(nums []int, k int) int {
    // partitionK
    return partitionK(nums, k)
}

func partitionK(nums []int, k int) int {
	var n int = len(nums)
	var quickSort func([]int, int, int)
	quickSort = func(nums []int, lo, hi int) {
		if lo < hi {
                mid := partition1(nums, lo, hi)
                if mid > n-k {
                        hi = mid - 1
                        quickSort(nums, lo, hi)
                } else if mid < n-k {
                        lo = mid + 1
                        quickSort(nums, lo, hi)
                } else {
                        return
                }
		}
	}
	quickSort(nums, 0, len(nums)-1)
	return nums[len(nums)-k]
}

func partition(nums []int, lo, hi int) int {
	var i, j, pivot int
	pivot = nums[lo]
	i, j = lo, hi
	for i < j {
        for i < j && nums[j] > pivot {
                j--
        }
        if i < j {
                nums[i], nums[j] = nums[j], nums[i]
                i++
        }
        for i < j && nums[i] < pivot {
                i++
        }
        if i < j {
                nums[i], nums[j] = nums[j], nums[i]
                j--
        }
	}
	return i
}

func partition1(nums []int, lo, hi int) int {
	var mid, pivot, start int
	mid = (hi-lo)/2 + lo
	if nums[lo] > nums[mid] {
		nums[lo], nums[mid] = nums[mid], nums[lo]
	}
	if nums[mid] > nums[hi] {
		nums[mid], nums[hi] = nums[hi], nums[mid]
	}
	if nums[lo] > nums[hi] {
		nums[lo], nums[hi] = nums[hi], nums[lo]
	}
	pivot = nums[mid]
	nums[mid], nums[hi] = nums[hi], nums[mid]
	start = lo
	for i := lo; i < hi; i++ {
		if nums[i] <= pivot {
			nums[start], nums[i] = nums[i], nums[start]
			start++
		}
	}
	nums[hi], nums[start] = nums[start], nums[hi]
	return start
}

func partition2(nums []int, lo, hi int) int {
	rander := rand.New(rand.NewSource(time.Now().UnixNano()))
	randIdx := rander.Intn(hi-lo+1) + lo
	var pivot, start int
	pivot = nums[randIdx]
	nums[randIdx], nums[hi] = nums[hi], nums[randIdx]
	start = lo
	for i := lo; i < hi; i++ {
		if nums[i] < pivot {
			nums[start], nums[i] = nums[i], nums[start]
			start++
		}
	}
	nums[start], nums[hi] = nums[hi], nums[start]
	return start
}

func partition3(nums []int, lo, hi int) int {
	var pivot, i, j int
	pivot = nums[lo]
	i, j = lo, hi
	for i < j {
        for i < j && nums[j] >= pivot {
                j--
        }
        for i < j && nums[i] <= pivot {
                i++
        }
        nums[i], nums[j] = nums[j], nums[i]
	}
	nums[lo], nums[i] = nums[i], nums[lo]
	// nums[lo], nums[j+1] = nums[j+1], nums[lo]
	return i
}

上述 parition 算法中只有第一个强制移动的分区算法可以通过全部测试用例,其他只能通过 40 个(共 41 个测试用例)。

小根堆

func findKthLargest(nums []int, k int) int {
    h := myHeap{}
    heap.Init(&h)
    for i := 0; i < len(nums); i++ {
        heap.Push(&h, nums[i])
        if h.Len() > k {
            heap.Pop(&h)
        }
    }
    return heap.Pop(&h).(int)
}

type myHeap []int
func (h myHeap) Len() int {
    return len(h)
} 
func (h myHeap) Swap(i, j int) {
    h[i], h[j] = h[j], h[i]
} 
func (h myHeap) Less(i, j int) bool {
    return h[i] < h[j]
}
func (h *myHeap) Push(x any) {
    *h = append(*h, x.(int))
}
func (h *myHeap) Pop() any {
    x := (*h)[h.Len()-1]
    *h = (*h)[:h.Len()-1]
    return x
} 

图示

以后补