Merging k sorted lists

Chapter 9 built singly linked lists and moved their nodes between owners with take. Here you have k linked lists, each already sorted, and you want one sorted list that contains all their nodes.

Section 10.2 compares three orders of merging. Sections 10.3 to 10.6 implement them, and section 10.7 lists the time and memory each program needs. Appendix C keeps the first, much longer version of this chapter, with eleven programs.

Two letters appear throughout. k is the number of lists. N is the total number of nodes across all of them.

10.1 Merging two sorted lists

Every program in this chapter merges two sorted lists at a time. Compare the front nodes of the two lists. Move the smaller one to the end of the output. Repeat until one list is empty, then attach the other list whole. Figure 10.1 works through an example.

With linked lists, “move a node to the output” does not copy the value. It detaches the node from the front of

10.2 Three strategies for k lists

Figure 10.2 shows three ways to extend two-list merging to k lists.

One at a time. Merge list 2 into list 1, then list 3 into the result, and so on. Each merge walks the whole result so far, which keeps growing. The first nodes are walked k times, so the total is about O(kN).

Pairwise rounds. Merge the lists in pairs, like the rounds of a tournament. After the first round there are k / 2 lists, after the second k / 4, and so on. In every round, each node takes part in exactly one merge. There are about log₂ k rounds, so the total is O(N log k).

A heap of heads. Put the front node of each list into a min-heap, the structure from chapter 5. Pop the smallest, append it to the output, and push the next node from the same list. Each node goes through one push and one pop on a heap of at most k entries. Each costs O(log k), so the total is again O(N log k).

Animation 10.1 counts node moves for eight lists of three nodes: pairwise rounds against one list at a time.

For 100 lists of 10,000 nodes each, N is a million. O(kN) is about 100 million steps. O(N log k) is about 7 million.

The pairwise strategy has a compact in-place form. It keeps the lists in a Vec and merges positions that are a growing interval apart (figure 10.3). The code in section 10.3 calls the interval gap.

In each round, the list at position i absorbs the list at i + interval. The position i runs through 0, 2 × interval, 4 × interval, and so on. The odd list out, number 4 here, waits until the interval is large enough to reach it.

10.3 Merging with the lists swapped by value

The merge builds its output after a placeholder node, the dummy head, so the first real node needs no special case. tail always points at the last node of the output (figure 10.4).

merge_two takes both lists by value. When the right list has the smaller front node, the two lists swap, so left always holds the node to move. The only mutable reference is tail. This keeps the loop readable and reduces the mutable borrow fatigue: there is one borrow to follow.

Listing 10.1 The node type and merge_two (lines 4 to 27). src/bin/merge_k_sorted_lists_owned.rs

type Link = Option<Box<Node>>;

struct Node {
    val: i32,
    next: Link,
}

/// Merge two sorted lists. O(n) time, O(1) extra memory.
fn merge_two(mut left: Link, mut right: Link) -> Link {
    let mut dummy = Box::new(Node { val: 0, next: None });
    let mut tail = &mut dummy;

    while let (Some(l), Some(r)) = (&left, &right) {
        if r.val < l.val {
            (left, right) = (right, left); // left now holds the smaller front
        }
        let mut node = left.unwrap(); // safe: checked by while let
        left = node.next.take();
        tail = tail.next.insert(node);
    }

    tail.next = left.or(right);
    dummy.next
}

while let (Some(l), Some(r)) = (&left, &right) runs while both lists have a node. &left and &right are shared references, used only to read the two front values.

If the right front is smaller, (left, right) = (right, left) swaps the two lists. The assignment moves the two Option values; no node moves, and no reference is taken. After it, the smaller front node is always in left.

left.unwrap() moves that list out of left. It cannot fail, because the while let checked that left is Some. node.next.take() cuts the front node off, and the rest of the list goes back into left.

tail.next.insert(node) attaches the node at the end of the output and returns a mutable reference to it. That node is the new tail.

When the loop ends, one list is empty. left.or(right) returns the other one, and tail.next takes it whole.

Animation 10.2 runs merge_two one line at a time on the lists of the trace below. The table beside the lists shows what each variable owns or borrows after every line. The last part reads l after the swap, which does not compile.

Trace the merge of 1 → 4 with 2 → 3 → 5:

Each pass moves one node, so merging n nodes takes O(n) time. The merge uses O(1) extra memory, and allocates only the dummy node.

Listing 10.2 merge_k (lines 29 to 47).

/// Merge k sorted lists in place: round 1 merges lists 1 apart, round 2 lists
/// 2 apart, and so on. O(N log k) time, O(1) extra memory.
fn merge_k(mut lists: Vec<Link>) -> Link {
    if lists.is_empty() {
        return None;
    }

    let size = lists.len();
    let mut gap = 1;
    while gap < size {
        let mut i = 0;
        while i + gap < size {
            lists[i] = merge_two(lists[i].take(), lists[i + gap].take());
            i += gap * 2;
        }
        gap *= 2;
    }
    lists[0].take()
}

merge_k follows figure 10.3. The outer while doubles gap. The inner while steps i by gap * 2 and merges position i with position i + gap. lists[i].take() moves each list out of its slot and leaves None, so merge_two receives two owned lists.

Animation 10.3 runs the merge of the trace. The last part leaves out the line after the loop.

$ cargo run --bin merge_k_sorted_lists_owned
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
[1, 1, 1, 3, 5]
[]

The complete program is listing 10.5 at the end of the chapter.

10.4 Borrowing two elements of one Vec

merge_k moves both lists out of the vector before the call. A merge that takes two &mut lists instead cannot be called on two slots of the same vector:

merge_two(&mut lists[i], &mut lists[i + gap]) // does not compile

Both &mut borrows are borrows of the whole vector lists, and Rust allows only one mutable borrow of a value at a time. The compiler does not reason about the two indexes being different.

split_at_mut(i + gap) splits the vector into two slices that do not overlap: the part before position i + gap, and the rest. It returns a mutable borrow of each. In the long version of this chapter, listing C.1 merges this way:

                let (left, right) = lists.split_at_mut(i + interval);
                left[i] = Self::merge_two(&mut left[i], &mut right[0]);

Now left[i] and right[0] come from different slices, and borrowing both is allowed. Section 10.6 uses split_at_mut to split a slice of lists in half.

10.5 A heap of list heads

This is the third strategy from figure 10.2. The heap does not hold nodes. It holds pairs of (front value, list number), wrapped in Reverse to make it a min-heap, as in chapter 5. Storing the list number instead of the node keeps the nodes in their lists, and avoids needing an ordering on nodes.

Listing 10.3 merge_k_lists (lines 10 to 33). src/bin/merge_k_sorted_lists_heap.rs

fn merge_k_lists(mut lists: Vec<Link>) -> Link {
    let mut heap = BinaryHeap::new();
    for (i, list) in lists.iter().enumerate() {
        if let Some(node) = list {
            heap.push(Reverse((node.data, i)));
        }
    }

    let mut dummy = Box::new(ListNode {
        data: 0,
        next: None,
    });
    let mut tail = &mut dummy;

    while let Some(Reverse((_, i))) = heap.pop() {
        let mut node = lists[i].take().unwrap();
        lists[i] = node.next.take();
        if let Some(next) = &lists[i] {
            heap.push(Reverse((next.data, i)));
        }
        tail = tail.next.insert(node);
    }
    dummy.next
}

Each pass of the loop:

1. Pops the smallest pair. Its second field, i, says which list holds that value.
2. Takes the front node out of lists[i] and puts the rest of that list back.
3. If the list still has a node, pushes that node’s value with i.
4. Attaches the taken node to the output with tail.next.insert(node), as in listing 10.1.

When two fronts have equal values, the pairs compare by list number, so the lower-numbered list goes first.

The heap needs only the current front of each list. That makes this strategy the right one when the lists are too large to hold in memory. One example is k sorted files read line by line.

$ cargo run --bin merge_k_sorted_lists_heap
1 1 2 3 4 4 5 6

The complete program is listing 10.6 at the end of the chapter.

10.6 Recursion on halves

This version merges from the top down. It splits the lists into two halves, merges each half recursively, and merges the two results. Its node holds an i32 named val, and List is the alias for a whole list:

#[derive(Debug)]
struct ListNode {
    val: i32,
    next: Option<Box<ListNode>>,
}

type List = Option<Box<ListNode>>;

Listing 10.4 The merge and the recursion on a slice (lines 11 to 34). src/bin/merge_k_sorted_lists_halves.rs

fn merge(a: List, b: List) -> List {
    match (a, b) {
        (None, rest) | (rest, None) => rest,
        (Some(mut left), Some(mut right)) => {
            if left.val > right.val {
                mem::swap(&mut left, &mut right);
            }
            left.next = merge(left.next.take(), Some(right));
            Some(left)
        }
    }
}

fn merge_k_slice(lists: &mut [List]) -> List {
    match lists.len() {
        0 => None,
        1 => lists[0].take(),
        _ => {
            let mid = lists.len() / 2;
            let (left_half, right_half) = lists.split_at_mut(mid);
            merge(merge_k_slice(left_half), merge_k_slice(right_half))
        }
    }
}

The first arm of merge, (None, rest) | (rest, None) => rest, handles an empty list on either side: the other list is the answer. In the second arm, mem::swap makes left the list with the smaller front. Its node stays at the front, and its next becomes the merge of the rest of left with right. No node is allocated; the boxes move.

merge_k_slice works on a slice, &mut [List], and splits it with split_at_mut from section 10.4. Splitting a slice creates two views of the same memory, so nothing is allocated to split.

merge calls itself once for each node of its output. The last merge outputs all N nodes, so the stack reaches O(N) frames. In section 9.6, a recursive drop overflowed the 8 MiB main-thread stack between 260,000 and 270,000 nodes. For long lists, use the loop of section 10.3.

$ cargo run --bin merge_k_sorted_lists_halves
Input lists:
[1 -> 4 -> 5]
[1 -> 3 -> 4]
[2 -> 6]

Merged list:
[1 -> 1 -> 2 -> 3 -> 4 -> 4 -> 5 -> 6]

The complete program is listing 10.7 at the end of the chapter.

10.7 The three programs side by side

In the table, N is the total number of nodes and k the number of lists. The extra memory column counts memory beyond the input Vec and the nodes. The stack column is the deepest the call stack goes, in frames.

• Time. Each node moves about log₂ k times: O(N log k). The “one at a time” strategy of figure 10.2 would be O(kN).
• Extra memory. The heap holds one entry per list: O(k). The other two programs reuse the input Vec.
• Stack. The recursive merge makes one call per node of its output, so its stack reaches O(N) frames. The loop and the heap use a fixed number of frames.
• Allocations. The gap rounds allocate one dummy node per merge, k - 1 in total. The recursive merge allocates nothing.

10.8 The complete files

Listing 10.5 Gap rounds with the lists swapped by value. src/bin/merge_k_sorted_lists_owned.rs

//! Merge k sorted lists: the gap loop, and a merge that swaps the two lists by
//! value so that `left` always holds the smaller front node.

type Link = Option<Box<Node>>;

struct Node {
    val: i32,
    next: Link,
}

/// Merge two sorted lists. O(n) time, O(1) extra memory.
fn merge_two(mut left: Link, mut right: Link) -> Link {
    let mut dummy = Box::new(Node { val: 0, next: None });
    let mut tail = &mut dummy;

    while let (Some(l), Some(r)) = (&left, &right) {
        if r.val < l.val {
            (left, right) = (right, left); // left now holds the smaller front
        }
        let mut node = left.unwrap(); // safe: checked by while let
        left = node.next.take();
        tail = tail.next.insert(node);
    }

    tail.next = left.or(right);
    dummy.next
}

/// Merge k sorted lists in place: round 1 merges lists 1 apart, round 2 lists
/// 2 apart, and so on. O(N log k) time, O(1) extra memory.
fn merge_k(mut lists: Vec<Link>) -> Link {
    if lists.is_empty() {
        return None;
    }

    let size = lists.len();
    let mut gap = 1;
    while gap < size {
        let mut i = 0;
        while i + gap < size {
            lists[i] = merge_two(lists[i].take(), lists[i + gap].take());
            i += gap * 2;
        }
        gap *= 2;
    }
    lists[0].take()
}

fn build(values: &[i32]) -> Link {
    let mut head = None;
    for &val in values.iter().rev() {
        head = Some(Box::new(Node { val, next: head }));
    }
    head
}

fn to_vec(mut current: Link) -> Vec<i32> {
    let mut out = Vec::new();
    while let Some(node) = current {
        out.push(node.val);
        current = node.next;
    }
    out
}

fn main() {
    let lists = vec![
        build(&[1, 4, 7]),
        build(&[2, 5, 8]),
        build(&[0, 3, 6, 9]),
        None,
    ];
    println!("{:?}", to_vec(merge_k(lists))); // [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
    println!(
        "{:?}",
        to_vec(merge_k(vec![build(&[1, 1, 5]), build(&[1, 3])]))
    ); // [1, 1, 1, 3, 5]
    println!("{:?}", to_vec(merge_k(vec![]))); // []
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn merges_two_lists_keeping_duplicates() {
        let merged = merge_two(build(&[1, 1, 5]), build(&[1, 3]));
        assert_eq!(to_vec(merged), vec![1, 1, 1, 3, 5]);
    }

    #[test]
    fn merges_with_empty_lists() {
        assert_eq!(to_vec(merge_two(None, build(&[2, 4]))), vec![2, 4]);
        assert_eq!(to_vec(merge_two(build(&[2, 4]), None)), vec![2, 4]);
        assert!(merge_two(None, None).is_none());
        assert!(merge_k(vec![]).is_none());
        assert!(merge_k(vec![None, None]).is_none());
    }

    #[test]
    fn merges_an_odd_number_of_lists() {
        let five = vec![
            build(&[9]),
            build(&[1, 5]),
            build(&[3]),
            build(&[2, 8]),
            build(&[4]),
        ];
        assert_eq!(to_vec(merge_k(five)), vec![1, 2, 3, 4, 5, 8, 9]);
    }

    #[test]
    fn merges_many_long_lists() {
        let lists: Vec<Link> = (0..7)
            .map(|k| {
                build(
                    &(0..100)
                        .map(|i| i * 7 + k)
                        .collect::<Vec<_>>(),
                )
            })
            .collect();
        let expected: Vec<i32> = (0..700).collect();
        assert_eq!(to_vec(merge_k(lists)), expected);
    }
}

Listing 10.6 A heap of list heads. src/bin/merge_k_sorted_lists_heap.rs

use std::{cmp::Reverse, collections::BinaryHeap};

struct ListNode {
    data: u32,
    next: Option<Box<ListNode>>,
}

type Link = Option<Box<ListNode>>;

fn merge_k_lists(mut lists: Vec<Link>) -> Link {
    let mut heap = BinaryHeap::new();
    for (i, list) in lists.iter().enumerate() {
        if let Some(node) = list {
            heap.push(Reverse((node.data, i)));
        }
    }

    let mut dummy = Box::new(ListNode {
        data: 0,
        next: None,
    });
    let mut tail = &mut dummy;

    while let Some(Reverse((_, i))) = heap.pop() {
        let mut node = lists[i].take().unwrap();
        lists[i] = node.next.take();
        if let Some(next) = &lists[i] {
            heap.push(Reverse((next.data, i)));
        }
        tail = tail.next.insert(node);
    }
    dummy.next
}

fn main() {
    let mut lists = Vec::new();
    for values in [vec![1, 4, 5], vec![1, 3, 4], vec![2, 6]] {
        let mut head = None;
        for data in values.into_iter().rev() {
            head = Some(Box::new(ListNode { data, next: head }));
        }
        lists.push(head);
    }

    let mut cur = merge_k_lists(lists);
    while let Some(node) = cur {
        print!("{} ", node.data);
        cur = node.next;
    }
    println!();
}

Listing 10.7 Recursion on halves. src/bin/merge_k_sorted_lists_halves.rs

use std::mem;

#[derive(Debug)]
struct ListNode {
    val: i32,
    next: Option<Box<ListNode>>,
}

type List = Option<Box<ListNode>>;

fn merge(a: List, b: List) -> List {
    match (a, b) {
        (None, rest) | (rest, None) => rest,
        (Some(mut left), Some(mut right)) => {
            if left.val > right.val {
                mem::swap(&mut left, &mut right);
            }
            left.next = merge(left.next.take(), Some(right));
            Some(left)
        }
    }
}

fn merge_k_slice(lists: &mut [List]) -> List {
    match lists.len() {
        0 => None,
        1 => lists[0].take(),
        _ => {
            let mid = lists.len() / 2;
            let (left_half, right_half) = lists.split_at_mut(mid);
            merge(merge_k_slice(left_half), merge_k_slice(right_half))
        }
    }
}

fn create_list(arr: &[i32]) -> List {
    let mut head = None;
    let mut tail = &mut head;

    for &val in arr {
        *tail = Some(Box::new(ListNode { val, next: None }));
        if let Some(node) = tail {
            tail = &mut node.next;
        }
    }

    head
}

fn print_list(mut list: &List) {
    let mut elements = Vec::new();
    while let Some(node) = list {
        elements.push(node.val.to_string());
        list = &node.next;
    }
    println!("[{}]", elements.join(" -> "));
}

fn main() {
    let mut lists = vec![
        create_list(&[1, 4, 5]),
        create_list(&[1, 3, 4]),
        create_list(&[2, 6]),
    ];

    println!("Input lists:");
    for list in &lists {
        print_list(list);
    }

    let merged = merge_k_slice(&mut lists);

    println!("\nMerged list:");
    print_list(&merged);
}

Chapter 11 moves from lists to trees, where each node has up to two children. It ends with tries, where each node has one child per letter.

Exercises

1. Write the “one at a time” merge and time it against listing 10.5 on 200 lists of 5,000 nodes each.
2. Make merge_two from listing 10.1 generic over any T: Ord.
3. Write a k-way merge of sorted text files with the heap strategy, keeping only one line per file in memory.
4. Rewrite merge from listing 10.4 as a loop with the swap of listing 10.1. Test it on two lists of a million nodes each.