C语言实现数组堆并解决TopK问题

还是先定义结构体

typedef int HPDataType;
typedef struct {
    HPDataType* array;
    int size;
    int capacity;
} HP;

void HeapInit(HP* php) {
    assert(php);
    php->array = NULL;
    php->capacity = php->size = 0;
}

首先是它的初始化。

void HeapDestroy(HP* php) {
    assert(php);
    free(php->array);
    php->capacity = php->size = 0;
}

堆的销毁。

void HeapPrint(HP* php) {
    assert(php);
    for (int i = 0; i < php->size; i++) {
        printf("%d  ", php->array[i]);
    }
    printf("\n");
}

打印堆。

bool HeapIsEmpty(HP* php) {
    assert(php);
    return php->size == 0;
}

判断堆是否为空。

void HeapCheckCapacity(HP* php) {
    if (php->capacity == php->size) {
        int newCapacity = php->capacity == 0 ? 4 : php->capacity * 2;
        HPDataType* tmpArray = (HPDataType*)realloc(php->array, sizeof(HPDataType) * newCapacity);
        if (tmpArray == NULL) {  // 检查 realloc
            printf("realloc failed!\n");
            exit(-1);
        }
        php->capacity = newCapacity;
        php->array = tmpArray;
    }
}

检查堆是否需要扩容。如果已满（capacity=size），我们就需要重新设定capacity和为数组分配空间，这里的realloc函数用来为数组重新分配空间。

void Swap(HPDataType* px, HPDataType* py) {
    HPDataType tmp = *px;
    *px = *py;
    *py = tmp;
}

交换函数，下面会用到。

void MinAdjustUp(int* arr, int child) {
    assert(arr);
    int father = (child - 1) / 2;
    while (child > 0) {
        if (arr[child] < arr[father]) {
            Swap(&arr[child], &arr[father]);
            child = father;
            father = (child - 1) / 2;
        }
        else {
            break;
        }
    }
}

小堆向上调整。传入child，跟据其计算出father的值。child>0是为了确保避免调整操作超出数组，这里的child=father和father=(child-1)/2可以让操作继续向上执行。如果无需交换，就直接break跳出此次循环。

void HeapPush(HP* php, HPDataType x) {
    assert(php);
    HeapCheckCapacity(php);
    php->array[php->size] = x;
    php->size++;
    MinAdjustUp(php->array, php->size - 1);
}

新元素入小堆。将新元素先添加到数组末尾，然后调用小堆向上调整操作。

void MinAdjustDown(int* arr, int n, int father) {
    int child = (father * 2) + 1;
    while (child < n) {
        if (child + 1 < n && arr[child + 1] < arr[child])
            child = child + 1;
        if (arr[father] > arr[child]) {
            Swap(&arr[father], &arr[child]);
            father = child;
            child = (father * 2) + 1;
        }
        else {
            break;
        }
    }
}

小堆向下调整操作。传入father，计算出child，这里的child是左孩子。while内第一个if是为了找到左右孩子中最小的，用来与father交换，记得更新father和child的值，以便操作能继续进行。

void HeapPop(HP* php) {
    assert(php);
    assert(!HeapIsEmpty(php));
    Swap(&php->array[0], &php->array[php->size - 1]);
    php->size--;
    MinAdjustDown(php->array, php->size, 0);
}

删除小堆的顶点。先交换堆顶和堆尾数据，然后删除堆尾，再将堆顶元素进行小堆向下调整。

HPDataType HeapTop(HP* php) {
    assert(php);
    assert(!HeapIsEmpty(php));
    return php->array[0];
}

获得堆顶元素。

void printTopK(int* arr, int N, int K) {
    HP hp;
    HeapInit(&hp);
    for (int i = 0; i < K; i++)
        HeapPush(&hp, arr[i]);
    for (int i = K; i < N; i++) {
        if (arr[i] > HeapTop(&hp)) {
            HeapPop(&hp);
            HeapPush(&hp, arr[i]);
        }
    }
    HeapPrint(&hp);
    HeapDestroy(&hp);
}

这个函数可以在N个元素中找到前K个最大元素，即TopK问题。如上，先创建一个小堆，然后将其它N-K个元素与堆头比较，如果大于，就删除堆顶，再将新元素推进堆中，数字越大，下沉越深。

循环结束后打印出的小堆就是前K个最大的元素。

void testPrintTopK() {
    int arr[] = { 51, 334, 87, 150, 25, 65, 66, 34, 72 };
    int N = sizeof(arr) / sizeof(arr[0]);
    int K = 5;

    printf("Original array: ");
    for (int i = 0; i < N; i++) {
        printf("%d  ", arr[i]);
    }
    printf("\n");

    printTopK(arr, N, K);
}

测试函数。用来验证上面方法的正确性。