Merge n arrays

merge.c

#include <stdlib.h>
#include <stdio.h>

void swap(int32_t **a, int32_t **b) {
    int32_t *tmp = *a;
    *a = *b;
    *b = tmp;
}

/* In-place update of subtree at `i` - O(log(n)) */
void update_heap(int32_t **heads, int32_t **tails, uint32_t size, uint32_t i) {
    int32_t least = i;
    int32_t right = 2*i+2;
    int32_t left = 2*i+1;

    if (left < size && *(heads[left]) < *(heads[least]))
        least = left;
    if (right < size && *(heads[right]) < *(heads[least]))
        least = right;
    if (least != i) {
        swap(heads + i,heads + least);
        swap(tails + i,tails + least);
        update_heap(heads, tails, size, least);
    }
}

/* In-place initialization - O(n) */
void make_heap(int32_t **heads, int32_t **tails, uint32_t size) {
    uint32_t i;
    for (i = size / 2; i > 0; i--) {
        update_heap(heads, tails, size, i);
    }
    update_heap(heads, tails, size, 0);
}

/*
 * The first element of the array is the length of the array. For nested
 * arrays, the first element is the number of nested arrays and the first
 * nested array has the same format and begins at the 2nd element.
 *
 * Runtime of this implementation is O(n*log(m)), where n is the total number
 * of elements to be merged and m is the number of lists to merge.
 */
int32_t *mergemany(int32_t *lists) {
    uint32_t length = *lists;
    uint32_t size = 0;
    uint32_t i = 0;
    /* Because having two of these makes the syntax easier. */
    int32_t **heads = malloc(sizeof(int32_t *) * length);
    int32_t **tails = malloc(sizeof(int32_t *) * length);
    int32_t *ret;
    int32_t *curr = lists + 1;

    /* Set up and counting. */
    for (i = 0; i < length; i++) {
        heads[i] = curr + 1;
        tails[i] = curr + *curr;
        size += *curr;
        curr += *curr + 1;
    }

    ret = malloc(sizeof(int32_t) * (size + 1));
    *ret = size;
    curr = ret + 1;

    /* In-place initialization - O(length) */
    make_heap(heads, tails, length);

    /* O(size*log(length)) */
    while (curr < ret + size + 1) {
        /* heads[0] is the head of a minheap */
        *(curr++) = *(heads[0]++);
        if (heads[0] > tails[0]) {
            /* Heap delete */
            heads[0] = heads[--length];
            tails[0] = tails[length];
        }
        /* O(log(length)) */
        update_heap(heads, tails, length, 0);
    }

    free(heads);
    free(tails);

    return ret;
}

void print_list(int32_t *list) {
    uint32_t size = list[0];
    for (uint32_t i = 1; i <= size; i++) {
        printf("%d\n", list[i]);
    }
}

int main() {
    int32_t arrays[] = {
        4,
            5,  1,3,7,11,13,
            5,  0,5,6,9,14,
            6,  2,4,8,10,12,18,
            5,  15,16,17,19,20
    };
    print_list(mergemany(arrays));
}