Find K Smallest Elements

Find the k smallest elements from an unsorted array

Language Selection

Choose your preferred programming language

Showing: Python

Find K Smallest Elements

Problem Statement

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

Input/Output Specifications

Input:
- nums: Array of integers
- k: Integer representing number of smallest elements to find (1 <= k <= nums.length)
Output: Array containing k smallest 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: [1,2] or [2,1]
Explanation: The 2 smallest elements are 1 and 2.

Example 2:

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

Solution Approaches

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

Algorithm Explanation:

Use a max-heap of size k to maintain the k smallest 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_smallest_elements(nums, k):
    """
    Find k smallest elements using max-heap
    Time: O(n log k)
    Space: O(k)
    """
    if k == 0:
        return []
    
    max_heap = []
    
    for num in nums:
        if len(max_heap) < k:
            heapq.heappush(max_heap, -num)  # Negative for max-heap
        elif num < -max_heap[0]:  # num < largest in heap
            heapq.heapreplace(max_heap, -num)
    
    return [-x for x in max_heap]  # Convert back to positive

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

Java:

import java.util.*;

class Solution {
    /**
     * Find k smallest elements using max-heap
     * Time: O(n log k)
     * Space: O(k)
     */
    public int[] findKSmallestElements(int[] nums, int k) {
        if (k == 0) return new int[0];
        
        // Max-heap using reverse order comparator
        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(k, Collections.reverseOrder());
        
        for (int num : nums) {
            if (maxHeap.size() < k) {
                maxHeap.offer(num);
            } else if (num < maxHeap.peek()) {
                maxHeap.poll();
                maxHeap.offer(num);
            }
        }
        
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = maxHeap.poll();
        }
        
        return result;
    }
}

Go:

import (
    "container/heap"
)

// MaxHeap for maintaining k smallest elements
type MaxHeap []int

func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i] > h[j] } // Max-heap
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

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

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

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

JavaScript:

class MaxHeap {
    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 max = this.heap[0];
        this.heap[0] = this.heap.pop();
        this.bubbleDown();
        return max;
    }
    
    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 maxChildIdx = 2 * idx + 1;
            if (2 * idx + 2 < this.heap.length && 
                this.heap[2 * idx + 2] > this.heap[maxChildIdx]) {
                maxChildIdx = 2 * idx + 2;
            }
            
            if (this.heap[idx] >= this.heap[maxChildIdx]) break;
            [this.heap[idx], this.heap[maxChildIdx]] = [this.heap[maxChildIdx], this.heap[idx]];
            idx = maxChildIdx;
        }
    }
}

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

C#:

using System;
using System.Collections.Generic;

public class Solution {
    /// <summary>
    /// Find k smallest elements using max-heap (PriorityQueue)
    /// Time: O(n log k)
    /// Space: O(k)
    /// </summary>
    public int[] FindKSmallestElements(int[] nums, int k) {
        if (k == 0) return new int[0];
        
        // Max-heap using negative priority (higher value = higher priority)
        var maxHeap = new PriorityQueue<int, int>(Comparer<int>.Create((a, b) => b.CompareTo(a)));
        
        foreach (int num in nums) {
            if (maxHeap.Count < k) {
                maxHeap.Enqueue(num, num);
            } else if (num < maxHeap.Peek()) {
                maxHeap.Dequeue();
                maxHeap.Enqueue(num, num);
            }
        }
        
        int[] result = new int[k];
        for (int i = 0; i < k; i++) {
            result[i] = maxHeap.Dequeue();
        }
        
        return result;
    }
}

Approach 2: Sorting (Simple but Less Efficient)

Algorithm Explanation:

Sort the entire array
Return first k elements

Implementation:

Python:

def find_k_smallest_sorting(nums, k):
    """
    Find k smallest elements using sorting
    Time: O(n log n)
    Space: O(1) if in-place sorting allowed
    """
    nums.sort()
    return nums[:k]

Approach 3: Quickselect with Modification (Optimal Average Time)

Algorithm Explanation:

Use quickselect to find the kth smallest element
Collect all elements <= kth smallest
Return exactly k elements

Implementation:

Python:

import random

def find_k_smallest_quickselect(nums, k):
    """
    Find k smallest 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 smallest 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_smallest = quickselect(nums_copy, 0, len(nums_copy) - 1, k - 1)
    
    result = []
    for num in nums:
        if num <= kth_smallest and len(result) < k:
            result.append(num)
    
    return result

Key Insights

Max-Heap for K Smallest: Use max-heap to track k smallest elements efficiently
Heap Size Management: Maintain exactly k elements in heap
Element Comparison: Compare new elements with heap root (largest of k smallest)
Memory Efficiency: Heap approach uses O(k) space vs O(n) for sorting

Edge Cases

k = 0: Return empty array
k = n: Return entire array
Duplicate Elements: Handle multiple occurrences of kth smallest
Single Element: Array with one element
All Same Values: All elements are identical

Test Cases

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

test_find_k_smallest_elements()

Interview Tips

Discuss Heap Choice: Explain why max-heap for k smallest elements
Compare Approaches: Heap vs sorting vs quickselect trade-offs
Handle Edge Cases: Empty input, k=0, k=n scenarios
Space Optimization: Heap approach saves memory for large arrays
Built-in Functions: Mention language-specific optimized functions

Follow-up Questions

Kth Smallest Element: Return only the kth smallest value
Streaming Data: Handle continuous input of numbers
Find K Largest: Modify approach for largest elements
Memory Constraints: What if only O(1) extra space allowed?
Stability: Maintain relative order of equal elements

Common Mistakes

Wrong Heap Type: Using min-heap instead of max-heap
Boundary Conditions: Not handling k=0 or k=n properly
Duplicate Elements: Incorrect handling of equal values
Heap Size: Not maintaining exactly k elements
Comparison Logic: Wrong condition for heap replacement

Concept Explanations

Max-Heap for K Smallest: This is the key insight - to find k smallest elements efficiently, we maintain a max-heap of size k. The largest element in this heap represents the “boundary” - any new element smaller than this boundary should replace it.

Heap vs Sorting Trade-off: While sorting gives O(n log n) time, the heap approach gives O(n log k). When k « n, this is significantly more efficient. Additionally, heap uses O(k) space compared to potential O(n) space for sorting.

When to Apply: Use this pattern for top-k problems where k is much smaller than n, streaming data scenarios, or when memory is constrained. The heap approach is particularly effective for finding k smallest/largest elements in large datasets.