K-way Merge and Two Heaps

1. Merge k Sorted Lists Hard

Problem

You are given k sorted linked lists. Merge them into one sorted linked list and return its head.

The idea

Seed a min-heap with the head node of every non-empty list. Pop the smallest node, link it onto the result, and push that node’s Next. A dummy head node removes the “is the result empty yet” special case. The priority is (value, listIndex), so ties pop in a fixed order. The node itself is the element, and it is never compared.

Solution

/// <summary>A singly linked list node for this page's merge problems.</summary>
public class ListNode(int val = 0, ListNode? next = null)
{
    public int Val = val;            // default 0: only the dummy head uses it
    public ListNode? Next = next;

    /// <summary>Linked list holding the values, in order.</summary>
    /// <param name="values">Values to link. Empty input gives null.</param>
    /// <returns>The head node, or null.</returns>
    /// <example>ListNode.Build([1, 2]) returns 1 -&gt; 2.</example>
    public static ListNode? Build(int[] values)
    {
        ListNode? head = null;
        // Build from the back, so each new node points at the one built before.
        // i starts at Length - 1, the last index.
        for (int i = values.Length - 1; i >= 0; i--)
        {
            head = new ListNode(values[i], head);
        }
        return head;
    }

    /// <summary>The values of a linked list, as a List.</summary>
    /// <param name="head">First node, or null for an empty list.</param>
    /// <returns>The values in order.</returns>
    /// <example>ListNode.Values(ListNode.Build([1, 2])) returns [1, 2].</example>
    public static List<int> Values(ListNode? head)
    {
        var output = new List<int>();
        // Invariant: output holds every value before head.
        for (var node = head; node is not null; node = node.Next)
        {
            output.Add(node.Val);
        }
        return output;
    }
}

public static class MergeKLists
{
    /// <summary>Merge k sorted linked lists into one sorted linked list.</summary>
    /// <param name="lists">Heads of sorted linked lists. Any of them may be null.</param>
    /// <returns>The head of the merged list, or null if every list is empty.</returns>
    /// <example>Merge([Build([1, 4, 5]), Build([1, 3, 4]), Build([2, 6])]) gives
    /// 1 1 2 3 4 4 5 6.</example>
    public static ListNode? Merge(ListNode?[] lists)
    {
        // Element: the node. Priority: (value, list index). The index makes
        // equal values pop in list order. Nodes are never compared.
        var heap = new PriorityQueue<(ListNode Node, int List), (int, int)>();
        for (int index = 0; index < lists.Length; index++)
        {
            // Skip empty lists: they have no head.
            if (lists[index] is { } head)
            {
                heap.Enqueue((head, index), (head.Val, index));
            }
        }

        var dummy = new ListNode();      // placeholder before the real head
        var tail = dummy;                // last node of the merged list so far

        // Invariant: dummy.Next .. tail is sorted and holds every popped node.
        // heap holds the current head of every list not yet used up.
        while (heap.TryDequeue(out var entry, out _))   // _ : priority, not needed
        {
            tail.Next = entry.Node;
            tail = entry.Node;

            // Same index: the next node comes from the same list, and that list
            // has no other entry in the heap, so the priority stays unique.
            if (entry.Node.Next is { } next)
            {
                heap.Enqueue((next, entry.List), (next.Val, entry.List));
            }
        }
        return dummy.Next;               // skip the placeholder
    }
}

Walkthrough

Lists [1, 4, 5], [1, 3, 4], [2, 6]. Heap priorities shown as (value, index):

• Seed: (1,0) (1,1) (2,2). Pop (1,0), push (4,0). Out 1.
• Pop (1,1), push (3,1). Out 1 1. The tie on value 1 was settled by index.
• Pop (2,2), push (6,2). Then pop (3,1), push (4,1). Out 1 1 2 3.
• Pop (4,0), push (5,0). Pop (4,1), nothing to push. Out 1 1 2 3 4 4.
• Pop (5,0), then (6,2). Heap empty. Out 1 1 2 3 4 4 5 6.

Reading the figure. Each row is the state after one step. Heap entries show the priority (value, list index), smallest first. The amber entry is the next to pop. The violet entry was just pushed from the list that lost its head. Notice row 2: two entries tie on value 1, and the list index settles the order. Green cells are the merged output.

Edge cases to raise

• Empty input, or every list null: the heap starts empty and the result is null.
• One list: it comes back unchanged, at O(N log 1) = O(N).
• Duplicate values across lists: the tuple priority gives them a fixed order. Without it the output values are still right, but which node comes first is not fixed.
• The other O(N log k) answer is divide and conquer. Merge lists in pairs, halving the count each round, as in merge sort.

2. Kth Smallest Element in a Sorted Matrix Medium

Problem

An n × n matrix has every row and every column sorted in ascending order. Return the k-th smallest element, counting duplicates.

The idea

Each row is a sorted list, so this is a K-way merge over n rows that stops early. Seed the heap with the first cell of each row. Pop k - 1 times, pushing the cell to the right each time. The heap top is then the k-th smallest. Only the first min(n, k) rows can matter, because the k-th smallest cannot sit below row k.

Solution

public static class KthSmallestInMatrix
{
    /// <summary>The k-th smallest value in a matrix with sorted rows and columns.</summary>
    /// <param name="matrix">A non-empty n x n matrix. Rows and columns ascend.</param>
    /// <param name="k">Rank to find, 1 to n * n.</param>
    /// <returns>The k-th smallest value, duplicates counted.</returns>
    /// <exception cref="ArgumentOutOfRangeException">k is outside 1..n * n.</exception>
    /// <example>Find([[1, 5, 9], [10, 11, 13], [12, 13, 15]], 8) returns 13.</example>
    public static int Find(int[][] matrix, int k)
    {
        int n = matrix.Length;
        // 1 is the smallest rank. n * n is the count of cells, the largest rank.
        ArgumentOutOfRangeException.ThrowIfLessThan(k, 1);
        ArgumentOutOfRangeException.ThrowIfGreaterThan(k, n * n);

        // Seed with column 0 of each row. Element: (row, col). Priority: the value.
        // Rows past k cannot hold the answer, so Math.Min(n, k) rows is enough.
        var heap = new PriorityQueue<(int Row, int Col), int>();
        for (int row = 0; row < Math.Min(n, k); row++)
        {
            heap.Enqueue((row, 0), matrix[row][0]);   // 0: first column
        }

        // Pop k - 1 times: after that the smallest k - 1 are gone, and the top
        // of the heap is the k-th smallest.
        // Invariant: heap holds the next unused cell of every seeded row.
        for (int popped = 0; popped < k - 1; popped++)
        {
            var (row, col) = heap.Dequeue();
            // col + 1 is the next cell to the right. Push it if the row has one.
            if (col + 1 < n)
            {
                heap.Enqueue((row, col + 1), matrix[row][col + 1]);
            }
        }

        // The top's priority is its value. TryPeek returns both parts.
        heap.TryPeek(out _, out int answer);
        return answer;
    }
}

Walkthrough

Matrix rows [1, 5, 9], [10, 11, 13], [12, 13, 15], k = 8:

• Seed: 1, 10, 12, one per row.
• Pops 1 to 3 remove 1, 5, 9. Row 0 is used up, so pop 3 pushes nothing.
• Pops 4 and 5 remove 10, 11 and push row 1’s 11, then 13.
• Pop 6 removes 12 and pushes row 2’s 13. Pop 7 removes one 13.
• Seven pops done. The other 13 is on top. Answer 13.

Reading the figure. Each cell shows its value and the pop that removed it. Blue cells were popped. The order snakes across rows, because the heap always holds the next cell of each row. The green cell is on top of the heap after k - 1 = 7 pops. The grey 15 is never touched.

The two 13s tie, and PriorityQueue may pop either one first. The answer is the value 13 either way, so this problem does not need a tie-break.

Edge cases to raise

• k = 1: no pops. Answer matrix[0][0].
• k = n * n: every cell but one is popped. Answer matrix[n - 1][n - 1].
• Duplicates: each copy counts once, which the heap does for free.
• k out of range: the guard throws ArgumentOutOfRangeException. Without it, Dequeue on an empty queue throws InvalidOperationException.
• Large k: binary search on the value wins. Guess a value, count cells at or below it with a staircase walk in O(n), and narrow. That is O(n log(max - min)).

3. Find Median from Data Stream Hard

Problem

Design a class with AddNum(num), which adds a number from a stream, and FindMedian(), which returns the median of every number so far. Both should be fast for millions of calls.

The idea

Sorting on every call is O(n log n). Inserting into a sorted List<int> is O(n). Two heaps give O(log n) to add and O(1) to read. The lower half lives in a max-heap built with a reversed comparer. The upper half lives in a min-heap. Lower is allowed one extra item, so with an odd count the median is the lower top.

Solution

/// <summary>Running median with two heaps.</summary>
/// <remarks>lower is a max-heap holding the smaller half. upper is a min-heap holding
/// the larger half. lower.Count is upper.Count or upper.Count + 1, and every lower
/// value is at most every upper value.</remarks>
/// <example>Add 1 and 2, then FindMedian() returns 1.5. Add 3, then it returns 2.0.</example>
public class MedianFinder
{
    // Reversed comparer: b before a, so the largest value sits on top.
    private readonly PriorityQueue<int, int> lower =
        new(Comparer<int>.Create((a, b) => b.CompareTo(a)));
    private readonly PriorityQueue<int, int> upper = new();   // plain min-heap

    /// <summary>Add num to the stream in O(log n).</summary>
    /// <param name="num">The new number.</param>
    /// <example>finder.AddNum(5) stores 5.</example>
    public void AddNum(int num)
    {
        // Push into lower, then move lower's largest to upper. This keeps every
        // lower value at most every upper value, whatever num was.
        lower.Enqueue(num, num);
        int moved = lower.Dequeue();
        upper.Enqueue(moved, moved);

        // Upper may now hold one more than lower. Move its smallest back,
        // so lower is equal or exactly 1 bigger.
        if (upper.Count > lower.Count)
        {
            int back = upper.Dequeue();
            lower.Enqueue(back, back);
        }
    }

    /// <summary>Median of every number added so far, in O(1).</summary>
    /// <returns>The median as a double.</returns>
    /// <exception cref="InvalidOperationException">No number has been added.</exception>
    /// <example>After adding 1 and 2, FindMedian() returns 1.5.</example>
    public double FindMedian()
    {
        // 0: lower is empty only when both are, because lower is never smaller.
        if (lower.Count == 0)
        {
            throw new InvalidOperationException("no numbers added yet");
        }

        // Odd count: lower holds the extra one, and its top is the median.
        if (lower.Count > upper.Count)
        {
            return lower.Peek();
        }

        // Even count: average the two middle values. long stops overflow on the
        // sum. / 2.0 is floating division, so 1 and 2 give 1.5, not 1.
        return ((long)lower.Peek() + upper.Peek()) / 2.0;
    }
}

Walkthrough

Add 1, 2, 3. Heaps shown as plain values:

• Add 1: lower {1} → move 1 up → upper too big, move back. lower {1}, upper {}. Median 1.0.
• Add 2: lower {1, 2} → move max 2 up. lower {1}, upper {2}. Median (1 + 2) / 2.0 = 1.5.
• Add 3: lower {1, 3} → move 3 up → upper {2, 3} too big, move 2 back. lower {1, 2}, upper {3}. Median 2.0.

Reading the figure. Each column is the state after one AddNum. Blue cells are the lower half and violet the upper half, shown sorted. The amber cells are the two heap tops, the only values FindMedian reads. With an odd count, lower is one bigger and its top is the answer.

Edge cases to raise

• FindMedian before any add: throw, or return double.NaN. Ask which they want.
• Duplicates and negatives: no special case. The reversed comparer handles int.MinValue, where negation would not.
• Huge values: int.MaxValue plus int.MaxValue overflows an int. The long cast keeps the average right.
• Follow-up, all numbers in 0..100: keep 101 counters and walk them. O(1) add, O(100) read.
• Follow-up, a sliding window median: the same two heaps plus lazy deletion, since PriorityQueue cannot remove one entry. Count removed values in a Dictionary<int, int> and skip them when they reach a top. A SortedSet of (value, index) is the other option.

4. Smallest Range Covering Elements from K Lists Hard

Problem

You have k sorted lists of integers. Find the smallest range [a, b] that includes at least one number from every list. A range is smaller if b - a is smaller, or if the widths tie and a is smaller.

The idea

Pick one number from each list. The range they cover runs from their min to their max. To shrink it, the only useful move is to raise the min, because raising anything else can only widen it. So keep a heap of the current pick from each list, which gives the min in O(1), and track the max by hand. Pop the min, record the range if it is the best yet, and replace it with the next number from the same list. Stop when any list runs out, because then no pick covers every list.

Solution

public static class SmallestRange
{
    /// <summary>Smallest [lo, hi] holding at least one value from every list.</summary>
    /// <param name="nums">k non-empty arrays, each sorted ascending.</param>
    /// <returns>[lo, hi] for the smallest such range. Ties go to the smaller lo.</returns>
    /// <exception cref="ArgumentException">nums is empty, or any array is empty.</exception>
    /// <example>Find([[4, 10, 15, 24, 26], [0, 9, 12, 20], [5, 18, 22, 30]]) returns
    /// [20, 24].</example>
    public static int[] Find(int[][] nums)
    {
        // 0: an empty input, or an empty array, leaves nothing to cover it with.
        if (nums.Length == 0 || nums.Any(row => row.Length == 0))
        {
            throw new ArgumentException("need at least one list, and no empty lists");
        }

        // One pick per array. Element: (array index, position). Priority: the value.
        // Position 0 is each array's smallest value.
        var heap = new PriorityQueue<(int List, int Pos), int>();
        int currentMax = int.MinValue;   // int.MinValue: any real value beats it
        for (int index = 0; index < nums.Length; index++)
        {
            heap.Enqueue((index, 0), nums[index][0]);
            // The heap gives the min. The max has to be tracked by hand.
            currentMax = Math.Max(currentMax, nums[index][0]);
        }

        // long.MaxValue: wider than any real range, so the first range always wins.
        // Widths are long because max - min of two ints can overflow an int.
        long bestWidth = long.MaxValue;
        int bestLo = 0, bestHi = 0;      // 0, 0: placeholders, set on the first pop

        // Invariant: heap holds exactly one value from every array, and currentMax
        // is the largest of them. So [heap top, currentMax] covers every array.
        while (true)
        {
            heap.TryDequeue(out var at, out int low);

            // Strict <: on a tie keep the earlier range, whose lo is smaller,
            // because lo only grows as the loop runs.
            long width = (long)currentMax - low;
            if (width < bestWidth)
            {
                (bestWidth, bestLo, bestHi) = (width, low, currentMax);
            }

            // Pos + 1 == Length: this array has no next value, so no later pick can
            // cover it. Every remaining range has been tried.
            int next = at.Pos + 1;
            if (next == nums[at.List].Length)
            {
                return [bestLo, bestHi];
            }

            int value = nums[at.List][next];
            heap.Enqueue((at.List, next), value);
            currentMax = Math.Max(currentMax, value);   // the new pick may be the max
        }
    }
}

Walkthrough

Lists [4, 10, 15, 24, 26], [0, 9, 12, 20], [5, 18, 22, 30]. Picks shown as a set:

• {4, 0, 5}, range [0, 5], width 5. Best so far. Replace 0 with 9.
• {4, 9, 5}, range [4, 9], width 5. A tie, so keep [0, 5]. Replace 4 with 10.
• Next ranges: [5, 10], [9, 18], [10, 18], [12, 18], [15, 20], [18, 24]. None is narrower than 5.
• {24, 20, 22}: pop 20, range [20, 24], width 4. New best.
• 20 was the last value of the second list, so stop. Answer [20, 24].

Reading the figure. Each bar is the range of the current picks, one row per step, top to bottom. The white dots are the three picks, one from each list. The left dot is always the heap’s min, and it is the one replaced next. Green bars are a new best. The amber bar ties on width and loses on the left end. The search stops after the last row, because list 1 has run out.

Edge cases to raise

• One list: the best range is a single value, [x, x], found on the first pop.
• The same value in every list: width 0, as in [[1, 2, 3], [1, 2, 3]], which gives [1, 1].
• Empty list in the input: no range can cover it. The guard throws ArgumentException. Ask what the interviewer expects.
• Extreme values: int.MaxValue - int.MinValue overflows an int. The width is a long for that reason.
• Why stop when one list ends: every later pick drops that list. The min can only rise from there, so no better covering range exists.