HeapProblem 4 of 18

k largest element in an array

Brute Force
Time: O(n log n)
Space: O(1)
Optimal
Time: O(n + k log n)
Space: O(n)

Problem Statement

Given an array of integers, find the K largest elements in the array.

Example:

  • Input: arr = [1, 23, 12, 9, 30, 2, 50], k = 3
  • Output: [50, 30, 23]
  • Explanation: The 3 largest elements are 50, 30, and 23.

Approach 1: Sorting (Brute Force)

Intuition

Sort the array in descending order and return the first K elements.

Algorithm

  1. Sort the array in descending order
  2. Return the first K elements
java
import java.util.Arrays;
import java.util.Collections;

public class Solution {
    public int[] kLargestElements(int[] arr, int k) {
        Integer[] boxed = Arrays.stream(arr).boxed().toArray(Integer[]::new);
        Arrays.sort(boxed, Collections.reverseOrder());
        
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = boxed[i];
        }
        return result;
    }
}

Complexity Analysis

  • Time Complexity: O(n log n) - Due to sorting
  • Space Complexity: O(1) - In-place sorting (ignoring output array)

Approach 2: Using Min-Heap of Size K

Intuition

Maintain a min-heap of size K. The heap always contains the K largest elements seen so far. For each new element, if it's larger than the heap's minimum, replace the minimum.

Algorithm

  1. Create a min-heap with first K elements
  2. For remaining elements:
    • If element > heap top, pop top and push element
  3. Heap contains K largest elements
java
import java.util.PriorityQueue;
import java.util.ArrayList;
import java.util.Collections;

public class Solution {
    public int[] kLargestElements(int[] arr, int k) {
        // Min-heap
        PriorityQueue<Integer> minHeap = new PriorityQueue<>();
        
        for (int num : arr) {
            minHeap.offer(num);
            
            if (minHeap.size() > k) {
                minHeap.poll();
            }
        }
        
        // Extract elements
        int[] result = new int[k];
        for (int i = k - 1; i >= 0; i--) {
            result[i] = minHeap.poll();
        }
        
        return result;
    }
}

Complexity Analysis

  • Time Complexity: O(n log k) - Each insertion/deletion in heap of size k is O(log k)
  • Space Complexity: O(k) - Heap stores k elements

Approach 3: Using Max-Heap (Optimal for small K)

Intuition

Build a max-heap from the entire array, then extract K elements.

Algorithm

  1. Build max-heap from array - O(n)
  2. Extract maximum K times - O(k log n)
java
import java.util.PriorityQueue;
import java.util.Collections;

public class Solution {
    public int[] kLargestElements(int[] arr, int k) {
        // Max-heap
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
        
        for (int num : arr) {
            maxHeap.offer(num);
        }
        
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = maxHeap.poll();
        }
        
        return result;
    }
}

Complexity Analysis

  • Time Complexity: O(n + k log n) - Building heap O(n), k extractions O(k log n)
  • Space Complexity: O(n) - Heap contains all elements

Approach 4: QuickSelect Variation

Intuition

Use QuickSelect to partition array such that the K largest elements are in the last K positions.

Algorithm

  1. Use QuickSelect to find the (n-k)th smallest element
  2. All elements after this pivot are the K largest
java
import java.util.Arrays;
import java.util.Random;

public class Solution {
    private Random rand = new Random();
    
    private int partition(int[] arr, int low, int high) {
        // Random pivot
        int randIdx = low + rand.nextInt(high - low + 1);
        int temp = arr[randIdx];
        arr[randIdx] = arr[high];
        arr[high] = temp;
        
        int pivot = arr[high];
        int i = low - 1;
        
        for (int j = low; j < high; j++) {
            if (arr[j] <= pivot) {
                i++;
                temp = arr[i];
                arr[i] = arr[j];
                arr[j] = temp;
            }
        }
        
        temp = arr[i + 1];
        arr[i + 1] = arr[high];
        arr[high] = temp;
        
        return i + 1;
    }
    
    private void quickSelect(int[] arr, int low, int high, int k) {
        if (low >= high) return;
        
        int pivotIdx = partition(arr, low, high);
        
        if (pivotIdx == k) {
            return;
        } else if (pivotIdx > k) {
            quickSelect(arr, low, pivotIdx - 1, k);
        } else {
            quickSelect(arr, pivotIdx + 1, high, k);
        }
    }
    
    public int[] kLargestElements(int[] arr, int k) {
        int n = arr.length;
        int[] arrCopy = arr.clone();
        
        // Find position n-k
        quickSelect(arrCopy, 0, n - 1, n - k);
        
        // Extract k largest
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = arrCopy[n - k + i];
        }
        
        // Sort descending
        Arrays.sort(result);
        for (int i = 0; i < k / 2; i++) {
            int temp = result[i];
            result[i] = result[k - 1 - i];
            result[k - 1 - i] = temp;
        }
        
        return result;
    }
}

Complexity Analysis

  • Time Complexity: O(n) average, O(n²) worst case
  • Space Complexity: O(1) - In-place (O(log n) recursion stack)

Key Takeaways

  1. Sorting is simplest but not optimal for small K relative to N
  2. Min-heap of size K is best when K << N: O(n log k) time, O(k) space
  3. Max-heap is better when K is close to N: O(n + k log n) time
  4. QuickSelect gives O(n) average but O(n²) worst case
  5. For K largest, use min-heap; for K smallest, use max-heap
  6. Python's heapq.nlargest(k, arr) uses an optimized heap algorithm
  7. Choice depends on K/N ratio and whether sorted output is needed
Previous
Maximum of all subarrays of size k
Next
Kth smallest and largest element in an unsorted array