Mdr-C-Tutorial

主题

30009：最小的前 k 个元素 ​

题目 ​

从不可修改的整数数组中选出最小的前 k 个元素，并按非递减顺序写入输出数组。

解析 ​

维护一个容量为 k 的最大堆。堆中保存目前见过的最小 k 个值；新值比堆顶更小时替换堆顶。扫描结束后，最大堆进行原地堆排序即可得到升序结果。输出只在所有分配和计算完成后写入，因此失败时输出数组不变。

答案 ​

c
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>

static void swap_i64(int64_t *left, int64_t *right) {
    int64_t temporary = *left;
    *left = *right;
    *right = temporary;
}

static void max_sift_down(
    int64_t *heap,
    size_t size,
    size_t index
) {
    for (;;) {
        size_t largest = index;
        if (index < size / 2) {
            size_t left = index * 2 + 1;
            size_t right = left + 1;
            if (heap[left] > heap[largest]) {
                largest = left;
            }
            if (right < size && heap[right] > heap[largest]) {
                largest = right;
            }
        }
        if (largest == index) {
            return;
        }
        swap_i64(&heap[index], &heap[largest]);
        index = largest;
    }
}

static void max_sift_up(int64_t *heap, size_t index) {
    while (index > 0) {
        size_t parent = (index - 1) / 2;
        if (heap[parent] >= heap[index]) {
            return;
        }
        swap_i64(&heap[parent], &heap[index]);
        index = parent;
    }
}

static void sort_heap_ascending(int64_t *heap, size_t size) {
    while (size > 1) {
        swap_i64(&heap[0], &heap[size - 1]);
        --size;
        max_sift_down(heap, size, 0);
    }
}

bool smallest_k_i64(
    const int64_t *items,
    size_t count,
    size_t k,
    int64_t *out_items
) {
    if (k > count || (count > 0 && items == NULL) ||
        (k > 0 && out_items == NULL)) {
        return false;
    }
    if (k == 0) {
        return true;
    }
    if (k > SIZE_MAX / sizeof(int64_t)) {
        return false;
    }
    int64_t *heap = malloc(k * sizeof *heap);
    if (heap == NULL) {
        return false;
    }
    size_t heap_size = 0;
    for (size_t i = 0; i < count; ++i) {
        if (heap_size < k) {
            heap[heap_size] = items[i];
            ++heap_size;
            max_sift_up(heap, heap_size - 1);
        } else if (items[i] < heap[0]) {
            heap[0] = items[i];
            max_sift_down(heap, heap_size, 0);
        }
    }
    sort_heap_ascending(heap, heap_size);
    for (size_t i = 0; i < k; ++i) {
        out_items[i] = heap[i];
    }
    free(heap);
    return true;
}

每个输入元素至多进行一次 O(log⁡k) 的堆调整，所以时间复杂度为 O(nlog⁡k)，额外空间为 O(k)。