The Algorithm Design Manual (3rd) Exercise Solution 'E4-18'

Table of Contents

Introduction
Repository
Exercise 4-18
- Question
- Solution
  - Outline
  - Code

Introduction

This series is my solutions for algorithm exercises in'The Algorithm Design Manual' 3rd edition.
It's my solutions, not model answers, so if you find some mistakes or more sophisticated solutions, please post in comment area below.

Repository

Here

Exercise 4-18

Question

Give an $O(n log k)$ -time algorithm that merges $k$ sorted lists with a total of $n$ elements into one sorted list. (Hint: use a heap to speed up the obvious $O(kn)$ -time algorithm).

Solution

Outline

Construct max-heap using the last element of each list as key.

When examining and extracting mininum, extract the last element and the list shrink by 1.
If the list is not empty, do bubble-down by its new last value.
Else, swap with the right-most element of the heap and do bubble-down using it.

Note: The assingment order to newly created merged array is reverse.

Code

/*
Construct max-heap using the last element of each list as key.

When examining and extracting mininum, extract the last element and
the list shrink by 1.
If the list is not empty, do bubble-down by its new last value.
Else, swap with the right-most element of the heap and do bubble-down
using it.

Note: The assingment order to newly created merged array is reverse.
*/

#include <algorithm>
#include <cassert>
#include <stdexcept>
#include <vector>
size_t pq_parent(size_t i) { return i / 2; }
size_t pq_young_child(size_t i) { return i * 2; }
size_t pq_top() { return 1; }
size_t pq_len(size_t x) { return x - 1; }

void pq_swap(std::vector<std::vector<int>>& q, size_t i1, size_t i2) {
    auto t{std::move(q[i1])};
    q[i1] = std::move(q[i2]);
    q[i2] = std::move(t);
}

void bubble_down(std::vector<std::vector<int>>& q, size_t p) {
    auto c{pq_young_child(p)};
    auto max_idx{p};
    for (size_t i = 0; i < 2; i++) {
        if (c + i <= pq_len(q.size()) && q[c + i].back() > q[max_idx].back())
            max_idx = c + i;
    }
    if (max_idx == p) return;
    // swap
    pq_swap(q, p, max_idx);
    // recurse
    bubble_down(q, max_idx);
}

std::vector<std::vector<int>> make_heap(
    const std::vector<std::vector<int>>& v) {
    std::vector<std::vector<int>> q(v.size() + 1);
    std::copy(v.cbegin(), v.cend(), q.begin() + 1);

    for (size_t i = pq_len(q.size()) / 2; i >= 1; i--) {
        bubble_down(q, i);
    }
    return q;
}

int examine_max(std::vector<std::vector<int>>& q) {
    if (pq_len(q.size()) <= 0) throw std::invalid_argument{"Empty queue"};
    auto v_max{q[pq_top()].back()};
    q[pq_top()].pop_back();
    if (q[pq_top()].empty()) {
        q[pq_top()] = std::move(q[pq_len(q.size())]);
        q.resize(q.size() - 1);
    }
    bubble_down(q, pq_top());
    return v_max;
}

std::vector<int> merge_sorted_lists(
    const std::vector<std::vector<int>>& lists) {
    auto q{make_heap(lists)};
    std::vector<int> result;
    // get total size
    size_t n{0};
    for (const auto& lst : lists) {
        n += lst.size();
    }
    // set size
    result.resize(n);

    // construct in reverse order
    for (size_t i = n; i-- > 0;) {
        result[i] = examine_max(q);
    }
    return result;
}

int main(int argc, char const* argv[]) {
    assert(merge_sorted_lists({{1}}) == (std::vector<int>{1}));
    assert(merge_sorted_lists({{2}, {1}}) == (std::vector<int>{1, 2}));
    assert(merge_sorted_lists({{2}, {1}, {-1}}) ==
           (std::vector<int>{-1, 1, 2}));

    assert(merge_sorted_lists({{1, 2}, {3, 4}, {5, 6}}) ==
           (std::vector<int>{1, 2, 3, 4, 5, 6}));
    assert(merge_sorted_lists({{1, 6}, {2, 4}, {3, 5}}) ==
           (std::vector<int>{1, 2, 3, 4, 5, 6}));

    assert(merge_sorted_lists({{1, 2, 3},
                               {30, 31, 32, 33, 34, 35},
                               {10, 11, 12, 13},
                               {20, 21, 22, 23, 24},
                               {-10, -9, -8, -7, -6}}) ==
           (std::vector<int>{-10, -9, -8, -7, -6, 1,  2,  3,  10, 11, 12, 13,
                             20,  21, 22, 23, 24, 30, 31, 32, 33, 34, 35}));
    assert(merge_sorted_lists({{1, 3, 10, 22},
                               {-9, 30, 31, 33, 34},
                               {2, 11, 13, 35},
                               {-10, -7, 20, 21, 23, 24},
                               {-8, -6, 12, 32}}) ==
           (std::vector<int>{-10, -9, -8, -7, -6, 1,  2,  3,  10, 11, 12, 13,
                             20,  21, 22, 23, 24, 30, 31, 32, 33, 34, 35}));
    return 0;
}