Find K Largest Elements

Find the k largest elements from an unsorted array

Language Selection

Choose your preferred programming language

Showing: Python

Find K Largest Elements

Problem Statement

Given an integer array nums and an integer k, return the k largest elements from the array. The result can be in any order.

Input/Output Specifications

Input:
- nums: Array of integers
- k: Integer representing number of largest elements to find (1 <= k <= nums.length)
Output: Array containing k largest elements in any order

Constraints

1 <= k <= nums.length <= 10^5
-10^4 <= nums[i] <= 10^4

Examples

Example 1:

Input: nums = [3,2,1,5,6,4], k = 2
Output: [5,6] or [6,5]
Explanation: The 2 largest elements are 5 and 6.

Example 2:

Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
Output: [4,5,5,6] (in any order)
Explanation: The 4 largest elements are 6, 5, 5, and 4.

Solution Approaches

Approach 1: Min-Heap (Priority Queue) - Optimal for Memory

Algorithm Explanation:

Use a min-heap of size k to maintain the k largest elements seen so far
Iterate through all elements:
- If heap size < k, add element
- If heap size = k and current element > heap top, remove top and add current
Return all elements in the heap

Implementation:

Python:

import heapq

def find_k_largest_elements(nums, k):
    """
    Find k largest elements using min-heap
    Time: O(n log k)
    Space: O(k)
    """
    if k == 0:
        return []
    
    min_heap = []
    
    for num in nums:
        if len(min_heap) < k:
            heapq.heappush(min_heap, num)
        elif num > min_heap[0]:
            heapq.heapreplace(min_heap, num)
    
    return list(min_heap)

# Alternative implementation using heapq.nlargest
def find_k_largest_elements_builtin(nums, k):
    """
    Find k largest elements using built-in function
    Time: O(n log k)
    Space: O(k)
    """
    return heapq.nlargest(k, nums)

Java:

import java.util.*;

class Solution {
    /**
     * Find k largest elements using min-heap
     * Time: O(n log k)
     * Space: O(k)
     */
    public int[] findKLargestElements(int[] nums, int k) {
        if (k == 0) return new int[0];
        
        PriorityQueue<Integer> minHeap = new PriorityQueue<>(k);
        
        for (int num : nums) {
            if (minHeap.size() < k) {
                minHeap.offer(num);
            } else if (num > minHeap.peek()) {
                minHeap.poll();
                minHeap.offer(num);
            }
        }
        
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = minHeap.poll();
        }
        
        return result;
    }
}

Go:

import (
    "container/heap"
)

// MinHeap for maintaining k largest elements
type MinHeap []int

func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

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

func (h *MinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

// findKLargestElements - Find k largest elements using min-heap
// Time: O(n log k), Space: O(k)
func findKLargestElements(nums []int, k int) []int {
    if k == 0 {
        return []int{}
    }
    
    minHeap := &MinHeap{}
    heap.Init(minHeap)
    
    for _, num := range nums {
        if minHeap.Len() < k {
            heap.Push(minHeap, num)
        } else if num > (*minHeap)[0] {
            heap.Pop(minHeap)
            heap.Push(minHeap, num)
        }
    }
    
    result := make([]int, k)
    for i := 0; i < k; i++ {
        result[i] = heap.Pop(minHeap).(int)
    }
    
    return result
}

JavaScript:

class MinHeap {
    constructor() {
        this.heap = [];
    }
    
    push(val) {
        this.heap.push(val);
        this.bubbleUp();
    }
    
    pop() {
        if (this.heap.length === 0) return null;
        if (this.heap.length === 1) return this.heap.pop();
        
        const min = this.heap[0];
        this.heap[0] = this.heap.pop();
        this.bubbleDown();
        return min;
    }
    
    peek() {
        return this.heap[0];
    }
    
    size() {
        return this.heap.length;
    }
    
    bubbleUp() {
        let idx = this.heap.length - 1;
        while (idx > 0) {
            const parentIdx = Math.floor((idx - 1) / 2);
            if (this.heap[parentIdx] <= this.heap[idx]) break;
            [this.heap[parentIdx], this.heap[idx]] = [this.heap[idx], this.heap[parentIdx]];
            idx = parentIdx;
        }
    }
    
    bubbleDown() {
        let idx = 0;
        while (2 * idx + 1 < this.heap.length) {
            let minChildIdx = 2 * idx + 1;
            if (2 * idx + 2 < this.heap.length && 
                this.heap[2 * idx + 2] < this.heap[minChildIdx]) {
                minChildIdx = 2 * idx + 2;
            }
            
            if (this.heap[idx] <= this.heap[minChildIdx]) break;
            [this.heap[idx], this.heap[minChildIdx]] = [this.heap[minChildIdx], this.heap[idx]];
            idx = minChildIdx;
        }
    }
}

/**
 * Find k largest elements using min-heap
 * Time: O(n log k)
 * Space: O(k)
 */
function findKLargestElements(nums, k) {
    if (k === 0) return [];
    
    const minHeap = new MinHeap();
    
    for (const num of nums) {
        if (minHeap.size() < k) {
            minHeap.push(num);
        } else if (num > minHeap.peek()) {
            minHeap.pop();
            minHeap.push(num);
        }
    }
    
    const result = [];
    while (minHeap.size() > 0) {
        result.push(minHeap.pop());
    }
    
    return result;
}

C#:

using System;
using System.Collections.Generic;

public class Solution {
    /// <summary>
    /// Find k largest elements using min-heap (PriorityQueue)
    /// Time: O(n log k)
    /// Space: O(k)
    /// </summary>
    public int[] FindKLargestElements(int[] nums, int k) {
        if (k == 0) return new int[0];
        
        var minHeap = new PriorityQueue<int, int>();
        
        foreach (int num in nums) {
            if (minHeap.Count < k) {
                minHeap.Enqueue(num, num);
            } else if (num > minHeap.Peek()) {
                minHeap.Dequeue();
                minHeap.Enqueue(num, num);
            }
        }
        
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = minHeap.Dequeue();
        }
        
        return result;
    }
}

Approach 2: Quickselect with Modification (Optimal Time)

Algorithm Explanation:

Use quickselect to find the kth largest element
Partition array around this pivot
Return all elements >= pivot value

Implementation:

Python:

import random

def find_k_largest_quickselect(nums, k):
    """
    Find k largest elements using quickselect
    Time: O(n) average, O(n²) worst
    Space: O(1)
    """
    def quickselect(arr, left, right, k):
        if left == right:
            return arr[left]
        
        pivot_index = random.randint(left, right)
        pivot_index = partition(arr, left, right, pivot_index)
        
        if k == pivot_index:
            return arr[k]
        elif k < pivot_index:
            return quickselect(arr, left, pivot_index - 1, k)
        else:
            return quickselect(arr, pivot_index + 1, right, k)
    
    def partition(arr, left, right, pivot_index):
        pivot_value = arr[pivot_index]
        arr[pivot_index], arr[right] = arr[right], arr[pivot_index]
        
        store_index = left
        for i in range(left, right):
            if arr[i] > pivot_value:  # For largest elements
                arr[store_index], arr[i] = arr[i], arr[store_index]
                store_index += 1
        
        arr[right], arr[store_index] = arr[store_index], arr[right]
        return store_index
    
    if not nums or k <= 0:
        return []
    
    nums_copy = nums[:]
    kth_largest = quickselect(nums_copy, 0, len(nums_copy) - 1, k - 1)
    
    result = []
    for num in nums:
        if num >= kth_largest:
            result.append(num)
    
    return result[:k]  # Handle duplicates

Key Insights

Min-Heap Approach: Maintains exactly k elements, memory efficient
Heap Operations: Push O(log k), Pop O(log k), Peek O(1)
Quickselect: Better average time complexity but unstable worst case
Built-in Functions: Many languages provide optimized implementations

Edge Cases

k = 0: Return empty array
k = n: Return entire array
Duplicate Elements: Handle multiple occurrences of kth largest
Single Element: Array with one element
Already Sorted: Array in ascending or descending order

Test Cases

def test_find_k_largest_elements():
    # Basic case
    assert sorted(find_k_largest_elements([3,2,1,5,6,4], 2)) == [5,6]
    
    # With duplicates
    assert sorted(find_k_largest_elements([3,2,3,1,2,4,5,5,6], 4)) == [4,5,5,6]
    
    # k = 1
    assert find_k_largest_elements([1], 1) == [1]
    
    # k = n
    result = find_k_largest_elements([3,1,2], 3)
    assert sorted(result) == [1,2,3]
    
    # All same elements
    result = find_k_largest_elements([5,5,5,5], 2)
    assert result == [5,5]
    
    # Negative numbers
    assert sorted(find_k_largest_elements([-1,-3,2,0,1], 3)) == [0,1,2]
    
    print("All tests passed!")

test_find_k_largest_elements()

Interview Tips

Ask About Order: Clarify if result order matters
Discuss Trade-offs: Heap vs sorting vs quickselect
Handle Duplicates: Explain how to handle multiple kth largest values
Memory Constraints: Min-heap is better for large arrays, small k
Follow-up Optimization: Mention built-in functions availability

Follow-up Questions

Find Kth Largest Element: Return only the kth largest value
Top K Frequent: Find k most frequent elements
Streaming Data: Handle continuous stream of numbers
Find K Smallest: Modify for smallest elements
Memory Limit: What if we can only use O(1) extra space?

Common Mistakes

Wrong Heap Type: Using max-heap instead of min-heap
Off-by-One Error: Confusing kth largest vs k largest elements
Duplicate Handling: Not considering elements with same value
Edge Case: Not handling k=0 or k=n properly
Modifying Original: Changing input array when not allowed

Concept Explanations

Min-Heap for K Largest: Counter-intuitive but efficient - we maintain the k largest elements seen so far in a min-heap. The smallest of these k elements (at heap root) can be quickly compared with new elements and replaced if needed.

When to Apply: Use this pattern for top-k problems, especially when k is much smaller than n. The heap approach is particularly useful for streaming data where elements arrive continuously.