215. Kth Largest Element in an Array

Problem

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.

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 assume k is always valid, 1 ≤ k ≤ array's length.

Problem Summary

Find the Kth largest element in an array. This problem is very classic. It can be implemented with O(n) time complexity.

Solution Approach

In the partition operation of quickselect, each partition operation returns a point, and the index of this pivot is the same as the index where this element would be in the final sorted array. Using this property, we can continuously divide the array interval and eventually find the Kth largest element. After performing one partition operation, if the index of this element is smaller than K, then continue performing partition operations in the interval on the right; if the index of this element is greater than K, then continue performing partition operations in the interval on the left; if they are equal, simply output the array element corresponding to this index.
The algorithm implemented with the quickselect idea has a time complexity of O(n) and a space complexity of O(logn). Since the proof process is very tedious, it will not be discussed here. For the specific proof, refer to Chapter 9, Section 2 of Introduction to Algorithms.

Code


package leetcode

import (
	"math/rand"
	"sort"
)

// Solution 1: Sorting; the sorting method is actually the fastest
func findKthLargest1(nums []int, k int) int {
	sort.Ints(nums)
	return nums[len(nums)-k]
}

// Solution 2: The theoretical basis of this method is that the index of the point obtained by partition is the index after the final sorting; based on this index, we can determine where the Kth largest number is
// Time complexity O(n), space complexity O(log n), worst-case time complexity O(n^2), space complexity O(n)
func findKthLargest(nums []int, k int) int {
	m := len(nums) - k + 1 // mth smallest, from 1..len(nums)
	return selectSmallest(nums, 0, len(nums)-1, m)
}

func selectSmallest(nums []int, l, r, i int) int {
	if l >= r {
		return nums[l]
	}
	q := partition(nums, l, r)
	k := q - l + 1
	if k == i {
		return nums[q]
	}
	if i < k {
		return selectSmallest(nums, l, q-1, i)
	} else {
		return selectSmallest(nums, q+1, r, i-k)
	}
}

func partition(nums []int, l, r int) int {
	k := l + rand.Intn(r-l+1) // This is an optimization, making the expected time complexity reduce to O(n), while the worst-case time complexity is O(n^2)
	nums[k], nums[r] = nums[r], nums[k]
	i := l - 1
	// nums[l..i] <= nums[r]
	// nums[i+1..j-1] > nums[r]
	for j := l; j < r; j++ {
		if nums[j] <= nums[r] {
			i++
			nums[i], nums[j] = nums[j], nums[i]
		}
	}
	nums[i+1], nums[r] = nums[r], nums[i+1]
	return i + 1
}

// Extension: Coding Interviews Offer 40. The smallest k numbers
func getLeastNumbers(arr []int, k int) []int {
	return selectSmallest1(arr, 0, len(arr)-1, k)[:k]
}

// It is implemented exactly the same as selectSmallest; only the return value no longer needs to be sliced, and nums can be returned directly
func selectSmallest1(nums []int, l, r, i int) []int {
	if l >= r {
		return nums
	}
	q := partition(nums, l, r)
	k := q - l + 1
	if k == i {
		return nums
	}
	if i < k {
		return selectSmallest1(nums, l, q-1, i)
	} else {
		return selectSmallest1(nums, q+1, r, i-k)
	}
}