Start learning free

LESSON 292 OF 348 · HEAP

Smallest K Integers in TypeScript: O(n log k) with a Max Heap

Practice O(n log k) Read the full explanation

Given an array of positive integers nums, return the smallest k values, in any order you want.

Example:

Input: nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4
Output: [1, 2, 2, 3]
Explanation: Smallest number is 1, 2nd smallest is 2, 
            3rd smallest is 2, 4th smallest is 3

The result can be in any order, [2, 1, 3, 2] is also a correct answer.

Note:

For this lesson, your algorithm should run in O(n log k) time and use O(k) extra space.
(There are faster solutions which we will discuss in future lessons)


STEP 1 OF 7
    Ctrl + Enter to check

    TYPESCRIPT · HEAP

    Smallest K in TypeScript, explained

    Problem overview

    The method kSmallest receives an array of positive integers, nums, and a number k. It returns the k smallest values, in any order.

    For nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5] and k = 4, the sorted array is [1, 2, 2, 3, 3, 5, 5, 6, 7, 9], so the four smallest values are 1, 2, 2 and 3. Duplicates count separately: both 2s are in the answer. [2, 1, 3, 2] is just as correct as [1, 2, 2, 3]. The statement asks for O(n log k) time and O(k) extra space.

    Approach: keep the k best so far in a max heap

    The obvious solution sorts the whole array and takes the first k values. That costs O(n log n) time and sorts far more than you need when k is small.

    The lesson's idea is to hold only k values at any moment. A heap is a structure where the largest (for a max heap) or smallest (for a min heap) value is always available at the top, and where insert and extract both cost O(log size). Here the heap is a max heap that stores the k smallest values seen so far. Its top is the largest of those k, which makes it the candidate for removal. When a new value arrives:

    • if it is smaller than the top, it deserves a place among the k smallest, and the old top no longer does;
    • if it is not smaller than the top, the heap already holds k values that are at least as good, so the newcomer is ignored.

    The heap never holds more than k values, so every operation costs O(log k), and n values make O(n log k). TypeScript has no built-in heap, so the lesson uses a ready-made max heap with the methods insert, extract and isEmpty.

    On the example with k = 4, the heap first takes 5, 9, 3, 6. The top is 9. Then 2 arrives: it is below 9, so 9 leaves and 2 enters (heap: 2, 3, 5, 6). Next 1 replaces 6, the next 3 replaces 5, and the next 2 replaces the other 3. The heap ends as 1, 2, 2, 3. The 7 and the last 5 never get in, because they aren't smaller than the top at that moment.

    Step by step

    1. Create the max heap smallestKNums. It will hold the k smallest values found so far.
    2. Loop over nums with an index i, from 0 up to but not including nums.length.
    3. Fill the heap first. While i < k, every value goes in with insert, because the heap isn't full yet.
    4. Once the heap is full, take the greatest value out with extract and store it in valueToAdd. One of two values, this one or nums[i], goes back in.
    5. Keep the smaller of the two. If nums[i] is below valueToAdd, replace valueToAdd with it. Otherwise the old top stays.
    6. Insert valueToAdd back, so the heap again holds exactly k values.
    7. Empty the heap into result. Declare result as number[], then repeat push(smallestKNums.extract()) while the heap isn't empty. The values come out from largest to smallest, which is fine because any order is accepted. Return result.

    The solution in TypeScript

    const algorithms = require("algorithms");
class Solution {
    kSmallest(nums: number[], k: number): number[] {
        let smallestKNums = new algorithms.DataStructures.Heap.MaxHeap();
        for (let i = 0; i < nums.length; ++i) {
            if (i < k) {
                smallestKNums.insert(nums[i]);
            } else {
                let valueToAdd = smallestKNums.extract();
                if (nums[i] < valueToAdd) {
                    valueToAdd = nums[i];
                }
                smallestKNums.insert(valueToAdd);
            }
        }

        let result: number[] = [];
        while (!smallestKNums.isEmpty()) {
            result.push(smallestKNums.extract());
        }
        return result;
    }
}

    The TypeScript details worth a closer look:

    • The heap comes from a package. The first line loads the algorithms module with require, and the heap class is reached through its nested names. You write no heap code in this lesson.
    • The empty array needs a type. let result: number[] = []; names the element type, so push is checked to receive numbers. An annotation like this exists only for the compiler and doesn't survive into the JavaScript.
    • Replacing without a second heap. Extracting the top and inserting the smaller of two values is a "replace the largest if the newcomer is smaller" operation, spelled out in two heap calls.

    Complexity

    Time: O(n log k). Each of the n values causes at most one extract and one insert on a heap of at most k values, each O(log k). Emptying the heap at the end takes k extracts, O(k log k), which is no more than the main loop because k is at most n in any meaningful call.

    Space: O(k). The heap holds at most k values. The returned array has k values too, which the problem requires anyway.

    Edge cases and common mistakes

    • Duplicates. With [2, 2, 2, 1] and k = 2, the answer is [1, 2] in some order. The test nums[i] < valueToAdd is strict, so an equal value changes nothing, and that is fine because the result looks the same either way.
    • k as large as the array. Every value is inserted in the first branch and the method returns all of them, from largest to smallest.
    • Using a min heap by mistake. A min heap would hand back the smallest value, which is exactly the one you want to keep. The point of the max heap is that its top is the value to throw out.
    • Forgetting the first k values. If the loop took the else path from the start, it would extract from an empty heap.
    • Expecting sorted output. The result comes out in descending order here. If you need ascending order, reverse it or sort it afterwards.
    • Sorting with the default comparator. If you solve this with sort() instead, remember it compares numbers as text: [10, 9, 1].sort() gives [1, 10, 9]. Pass (a, b) => a - b.

    Testing

    • nums = [5, 9, 3, 6, 2, 1, 3, 2, 7, 5], k = 4 gives [3, 2, 2, 1], which has the same values as [1, 2, 2, 3].
    • nums = [4, 1], k = 5 gives [4, 1].
    • nums = [2, 2, 2, 1], k = 2 gives the values 1 and 2.
    Keep going with TypeScriptAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free