215. Kth Largest Element in an Array

LeetCode 215. Kth Largest Element in an Array

Description

Given an integer array nums and an integer k, return the kth largest element in the array.

Note that it is the kth largest element in the sorted order, not the kth distinct element.

Example 1:

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

Constraints:

• 1 <= k <= nums.length <= 10^4

• -10^4 <= nums[i] <= 10^4

Tags

Divide and Conquer, Heap

Solution

Solution 1:

We can perform the partition operation adopted in quick sort algorithm until pivot == k.

• l == r, return nums[l];

• pivot == k, return nums[k];

• pivot > k, partition(left, pivot - 1);

• pivot < k, partition(pivot + 1, right);

Solution 2:

We use heap with capacity of k, where the root is the minimum element, to store the top-k largest elements in the given array. We keep pushing elements into the heap until the size of the heap is k. Then, we compare the current element with the root of the heap, if the root is smaller, we pop an element from the heap and push the current one into it. Finally, we return the root.

Complexity

Solution 1:

• Time complexity: $O(n^2)$, the worst case is we partition the array for n times. We can introduce random partition to speed up this up to $O(n)$ ;

• Space complexity: $O(\log(n))$, for recursion stack.

Solution 2:

• Time complexity: $O(n)$

• Space complexity: $O(k)$

Code

Solution 1 (Quick-sort):

func findKthLargest(nums []int, k int) int {
	if len(nums) == 0 {
		return 0
	}
	return topK(nums, 0, len(nums)-1, len(nums)-k)
}

func topK(data []int, left, right, k int) int {
	if left == right {
		return data[left]
	}
	pivot := partition(data, left, right)
	if pivot == k {
		return data[pivot]
	}
	if pivot > k {
		return topK(data, left, pivot-1, k)
	}
	return topK(data, pivot+1, right, k)
}

func partition(data []int, left, right int) int {
	pivotElem := data[right]
	i := left - 1
	for j := left; j < right; j++ {
		if data[j] < pivotElem {
			i++
			data[i], data[j] = data[j], data[i]
		}
	}
	i++
	data[i], data[right] = data[right], data[i]
	return i
}

Solution 2 (Heap):

type IntHeap []int

func (h IntHeap) Len() int {
	return len(h)
}

func (h IntHeap) Less(i, j int) bool {
	return h[i] < h[j]
}

func (h IntHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *IntHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
	x := (*h)[len(*h)-1]
	*h = (*h)[:len(*h)-1]
	return x
}

func findKthLargest1(nums []int, k int) int {
	h := make(IntHeap, 0, k)
	for _, num := range nums {
		if len(h) < k {
			heap.Push(&h, num)
		} else if num > h[0] {
			heap.Pop(&h)
			heap.Push(&h, num)
		}
	}
	return h[0]
}