Skip to main content
CodePeet

CodePeet 444 Complete DSA Sheet

Merging and Scheduling with a Heap

Editorial

One candidate from each list

Merging two sorted lists takes the smaller of the two fronts, again and again. With k lists you'd look at k fronts for every value, O(k) a step. Keep the fronts in a min-heap instead: its root is the smallest of them, so it's the smallest value left anywhere. Pop it, and push the next value from the list it came from. The heap never holds more than k values, one per list, and each step is O(log k):

A dashed zone titled the lists, sorted holds three lists, each in its own colour: list 1 is 2 6 9, list 2 is 1 3 8, list 3 is 4 5 7, each first value a solid tile and the rest pale. The heap holds those three first values as shaded discs in their lists' colours: 1 at the root, ringed in green, over 2 and 4. A green arrow takes 1 out into the merged row, with a pill reading 1 out first; a dashed arrow brings list 2's next value, 3, toward the root, with a pill reading 3 takes its place; the heap is labelled one per list. Focal line: pop the smallest, then push the next one from its list. Title: one candidate from each list.

Merge k sorted lists

The heap holds (value, list, position), so each popped value knows where its successor is:

// One candidate per list in a min-heap of (value, list, position): pop the smallest, push the next from its list.
vector<int> mergeKSorted(const vector<vector<int>>& lists) {
    priority_queue<tuple<int, int, int>, vector<tuple<int, int, int>>, greater<>> heap;
    for (int i = 0; i < (int)lists.size(); i++)
        if (!lists[i].empty()) heap.emplace(lists[i][0], i, 0);
    vector<int> out;
    while (!heap.empty()) {
        auto [v, i, j] = heap.top();
        heap.pop();
        out.push_back(v);
        if (j + 1 < (int)lists[i].size()) heap.emplace(lists[i][j + 1], i, j + 1);   // its list's next value
    }
    return out;
}

import heapq


def merge_k_sorted(lists):
    """One candidate per list in a min-heap of (value, list, position): pop the smallest, push the next from its list."""
    heap = [(l[0], i, 0) for i, l in enumerate(lists) if l]
    heapq.heapify(heap)
    out = []
    while heap:
        v, i, j = heapq.heappop(heap)
        out.append(v)
        if j + 1 < len(lists[i]):
            heapq.heappush(heap, (lists[i][j + 1], i, j + 1))   # its list's next value
    return out

// One candidate per list in a min-heap of {value, list, position}: pop the smallest, push the next from its list.
static List<Integer> mergeKSorted(int[][] lists) {
    PriorityQueue<int[]> heap = new PriorityQueue<>((x, y) -> Integer.compare(x[0], y[0]));
    for (int i = 0; i < lists.length; i++)
        if (lists[i].length > 0) heap.offer(new int[]{lists[i][0], i, 0});
    List<Integer> out = new ArrayList<>();
    while (!heap.isEmpty()) {
        int[] t = heap.poll();
        out.add(t[0]);
        int i = t[1], j = t[2];
        if (j + 1 < lists[i].length) heap.offer(new int[]{lists[i][j + 1], i, j + 1});   // its list's next value
    }
    return out;
}

// One candidate per list in Introduction to Heaps' MinHeap, ordered by value: [value, list, position].
function mergeKSorted(lists) {
    const heap = new MinHeap((x, y) => x[0] < y[0]);
    lists.forEach((l, i) => { if (l.length) heap.push([l[0], i, 0]); });
    const out = [];
    while (heap.a.length) {
        const [v, i, j] = heap.pop();
        out.push(v);
        if (j + 1 < lists[i].length) heap.push([lists[i][j + 1], i, j + 1]);   // its list's next value
    }
    return out;
}

Each printed [1, 1, 2, 3, 4, 4, 5, 6], [] and [], Merge k Sorted Lists' three examples as arrays. With linked lists the next value is node.next; in Python keep the index in the tuple too, as the first bug below shows. Each of the N values is pushed and popped once: O(N log k) time, O(k) extra memory. Step through it:

Merging [2, 6, 9], [1, 3, 8], [4, 5, 7] — The tree is the heap, one value from each list; the row is the merged output so far. Amber is the root, the next to go out; indigo, the value that just took its list's place.

Below, you run the merge. Pop the root, then click the list whose next value takes its place. A wrong click is refused with the reason.

Merge the Fronts: pop the smallest, refill from its list

Against merging one at a time

Two other ways come up in interviews: merge each list into everything before it, or merge neighbouring pairs, round after round, until one list is left. Timed in C++ on 1,000,000 values split into k sorted lists, the best of 3 runs:

ListsHeapOne list at a timePairs, round by round
1043.3 ms17.3 ms21.8 ms
10073.2 ms92.8 ms42.9 ms
1,00097.5 ms840.0 ms64.9 ms

One list at a time re-merges everything so far for every list, O(N·k): the fastest at k = 10, and 840 ms at k = 1,000. The heap and pairs are both O(N log k), and on arrays already in memory pairs win: each round is a run of plain comparisons, where the heap pays a push and a pop for every value. Sorting all 1,000,000 took 48 to 64 ms.

The heap's strengths are elsewhere. It holds one value per list, so it can merge lists too big to load at once, files or streams. And it can stop early: the 10 smallest of the 1,000 lists took 0.02 ms, which is how Design Twitter builds a feed from everyone you follow.

Greedy choices: the cheapest, the most, the next

A heap also serves greedy algorithms whose choices change as they go: take the best choice open right now, and new choices join the heap as they appear.

Minimum Cost Rope Connection: tying two ropes costs their total length, and the tied rope can be tied again. Always tie the two shortest:

import heapq


def min_cost_ropes(ropes):
    heap = list(ropes)
    heapq.heapify(heap)
    cost = 0
    while len(heap) > 1:
        a, b = heapq.heappop(heap), heapq.heappop(heap)   # the two shortest
        cost += a + b
        heapq.heappush(heap, a + b)                       # the tied rope goes back among the others
    return cost

For [4, 3, 2, 6] it returned 29. Tying them in the order given costs 31, and the two longest first 38. A rope's length is paid again every time the rope it's part of is tied, so the long ones should be tied last:

Two dashed zones, each a tree of ties for the ropes 4, 3, 2 and 6, the ropes drawn as discs, long ones orange and short ones teal, and each tie as a tile carrying its cost. Left, the two shortest first: 2 and 3 tie into 5, 5 and 4 into 9, 9 and 6 into 15; 2 and 3 sit three ties deep, marked paid ×3, 4 is paid ×2 and 6 ×1, and a green card with a tick reads 29 in all, 5 + 9 + 15. Right, the order given: 4 and 3 tie into 7, 7 and 2 into 9, 9 and 6 into 15; the long rope 4 sits three deep, paid ×3, and a red card with a cross reads 31 in all, 7 + 9 + 15. Focal line: every tie pays again for the ropes inside it. Title: tie the two shortest first.

Task Scheduler runs the task with the most copies left, then lets it cool for n slots before it can run again: a max-heap of counts and a queue of the tasks cooling down.

from collections import Counter, deque
import heapq


def least_interval(tasks, n):
    heap = [-c for c in Counter(tasks).values()]     # the most copies left on top
    heapq.heapify(heap)
    cooling = deque()                                # (when it may run again, copies left)
    time = 0
    while heap or cooling:
        time += 1
        if heap:
            left = heapq.heappop(heap) + 1           # run the task with the most left
            if left:
                cooling.append((time + n, left))
        if cooling and cooling[0][0] == time:
            heapq.heappush(heap, cooling.popleft()[1])
    return time

It returned 8, 6, 10 for the problem's three examples. A counting formula, max(len(tasks), (most − 1) · (n + 1) + tied), where tied counts the tasks with the most copies, gives the same totals (it agreed on 500 random cases), but only for this problem; the heap is what carries over.

Single-Threaded CPU takes tasks as they arrive and always runs the shortest one available. Sort the tasks by arrival; the heap holds the ones that have arrived, by (length, index):

import heapq


def get_order(tasks):
    order = sorted(range(len(tasks)), key=lambda i: tasks[i][0])   # by when each task arrives
    ready, out, time, k = [], [], 0, 0
    while len(out) < len(tasks):
        if not ready and time < tasks[order[k]][0]:
            time = tasks[order[k]][0]                               # idle: jump to the next arrival
        while k < len(order) and tasks[order[k]][0] <= time:
            i = order[k]
            heapq.heappush(ready, (tasks[i][1], i))                  # the shortest, then the lowest index
            k += 1
        length, i = heapq.heappop(ready)
        time += length
        out.append(i)
    return out

It returned [0, 2, 3, 1] and [4, 3, 2, 0, 1] for the examples, and agreed with a scan of every task on 500 random cases. When nothing is ready the clock jumps to the next arrival instead of ticking.

The same moves elsewhere: Hand of Straights starts every group at the smallest card left; Longest Happy String takes the letter with the most left, unless the last two are that letter, and then the next most.

Unlocking choices, and taking one back

IPO: with capital w you can start any project that needs at most w, and its profit adds to w. Sort the projects by the capital they need; as w grows, the ones you can now afford go into a max-heap by profit, and you take its top:

import heapq


def find_maximized_capital(k, w, profits, capital):
    projects = sorted(zip(capital, profits))     # by the capital each one needs
    affordable, i = [], 0                        # a max-heap of profits, as negatives
    for _ in range(k):
        while i < len(projects) and projects[i][0] <= w:
            heapq.heappush(affordable, -projects[i][1])
            i += 1
        if not affordable:
            break                                # nothing we can afford, now or later
        w -= heapq.heappop(affordable)           # the most profitable one we can afford
    return w

It returned 4 and 6 for the two examples.

Furthest Building You Can Reach: each climb takes bricks equal to its height, or one ladder. The ladders belong on the biggest climbs, but you only know which are biggest once you've seen them. So give every climb a ladder for now, in a min-heap; when there are more climbs than ladders, the smallest of them pays in bricks instead:

import heapq


def furthest_building(heights, bricks, ladders):
    climbs = []                                   # a min-heap of the climbs the ladders cover
    for i in range(len(heights) - 1):
        d = heights[i + 1] - heights[i]
        if d <= 0:
            continue
        heapq.heappush(climbs, d)
        if len(climbs) > ladders:
            bricks -= heapq.heappop(climbs)       # the smallest of them pays in bricks instead
            if bricks < 0:
                return i
    return len(heights) - 1
The buildings 4 12 2 7 3 18 20 3 19 drawn as buildings with lit windows, numbered 0 to 8, each with its height on a tag, and pills reading 10 bricks, 2 ladders and reaches building 7. Each climb is drawn in the gap before the taller building: wooden ladders on the climbs of 8 and 15, brick stacks on the climbs of 5 and 2, 7 of the 10 bricks. The climb of 16, from building 7 to 8, is a dashed red box with a cross: no ladder left and 3 bricks. Building 7 is green; building 8 is a dashed outline. Focal line: the ladders follow the biggest climbs so far. Title: the ladders go to the biggest climbs.

It returned 4, 7, 3 for the problem's three examples, in O(n log ladders): on [4, 12, 2, 7, 3, 18, 20, 3, 19] with 10 bricks and 2 ladders it reaches building 7.

Spotting it, and when it breaks

The clues. "k sorted lists", "k sorted arrays", "merge", the smallest across several sources; or a greedy that keeps taking the cheapest, the biggest or the earliest while new choices keep arriving: a heap of the choices open right now.

The near-misses. Two lists: a plain merge with two indexes. Every choice known at the start and none added: sort once (but a tied rope is a new choice, which is why the ropes need the heap). Only a count asked, with a fixed rule, as in Task Scheduler: a formula may do. Choices that only ever arrive in order: a queue.

When it breaks. A heap orders by one key. When the best choice depends on two (what you can afford, and what pays most), use one structure per condition, as IPO does. And a value that changes after it's pushed leaves a stale copy: push the new one and skip the old when it surfaces.

Bugs that cost a submission

Tuples whose ties compare the nodes (Python).

import heapq


class ListNode:
    def __init__(self, val, next=None):
        self.val, self.next = val, next


def merge_k_lists(lists):
    heap = [(node.val, node) for node in lists if node]
    heapq.heapify(heap)                          # two equal vals, and Python compares the nodes
    head = tail = ListNode(0)
    while heap:
        _, node = heapq.heappop(heap)
        tail.next = node
        tail = node
        if node.next:
            heapq.heappush(heap, (node.next.val, node.next))
    return head.next

On the example it raised TypeError: '<' not supported between instances of 'ListNode' and 'ListNode': two values tied at 1, so Python compared the next items, the nodes. Push (node.val, i, node) with a number no two entries share, such as the list's index. Java's and C++'s heaps compare only what the comparator compares.

A ladder whenever there's one left.

def furthest_building(heights, bricks, ladders):
    for i in range(len(heights) - 1):
        d = heights[i + 1] - heights[i]
        if d <= 0:
            continue
        if ladders:
            ladders -= 1                         # a ladder whenever one is left
        elif bricks >= d:
            bricks -= d
        else:
            return i
    return len(heights) - 1

On [4, 12, 2, 7, 3, 18, 20, 3, 19] with 10 bricks and 2 ladders it returned 4; the answer is 7. Its ladders went to the climbs of 8 and 5, and then 15 was too much for the bricks. The heap moves each ladder to a bigger climb when one turns up.

Ropes sorted once.

def min_cost_ropes(ropes):
    ropes = sorted(ropes)
    cost, tied = 0, ropes[0]
    for r in ropes[1:]:
        tied += r                                # the next shortest, tied onto everything so far
        cost += tied
    return cost

It returned 29 for [4, 3, 2, 6], the right answer, and 9 for [1, 1, 1, 1], where the heap's is 8. Each tied rope is a new rope, and it has to go back among the others: after 1 + 1, the next two shortest are the other two 1s, not the 2.

Check yourself

Quiz

1 / 6

Multiple choice

Merging k sorted lists with a heap, how many values does the heap hold at most?

Practice problems

In this order:

  1. Merge k Sorted Lists: one candidate per list.
  2. Minimum Cost Rope Connection: the two shortest, every time.
  3. Task Scheduler: the most copies left, then a cooldown.
  4. Hand of Straights: each group starts at the smallest card left.
  5. Single-Threaded CPU: the shortest task that has arrived.
  6. Longest Happy String: the most left, unless it makes three in a row.
  7. IPO: the best profit you can afford, as the capital grows.
  8. Furthest Building You Can Reach: ladders for the biggest climbs.
  9. Design Twitter: a merge of the feeds, stopping at 10.

Next, Two Heaps: the Running Median: a max-heap and a min-heap, half the values each.