【算法与数据结构】优先级队列 - 用二叉堆求数据流中的第K大的元素

数据流中的第K大的元素，总数据个数不足K个元素时返回-1。

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

typedef struct {
    int* data;
    int count;
    int capacity;
} KthLargest;

void heap_up(KthLargest* obj, int loc) {
    int p, c = loc;
    int tmp;
    while (c > 0) {
        p = (c - 1) >> 1;
        if (obj->data[c] < obj->data[p]) {
            tmp = obj->data[c];
            obj->data[c] = obj->data[p];
            obj->data[p] = tmp;
        }
        c = p;
    }
}
void heap_down(KthLargest* obj, int loc) {
    int p = loc;
    int bl, br, min;
    int tmp;
    while (((p<<1) + 1) < obj->count) {
        if (((p<<1) + 2) < obj->count) {
            bl = (p<<1) + 1;
            br = (p<<1) + 2;
            if (obj->data[bl] < obj->data[br]) {
                min = bl;
            } else {
                min = br;
            }
            if (obj->data[min] < obj->data[p]) {
                tmp = obj->data[min];
                obj->data[min] = obj->data[p];
                obj->data[p] = tmp;
                p = min;
            } else {
                break;
            }
        } else {
            bl = (p<<1) + 1;
            if (obj->data[bl] < obj->data[p]) {
                tmp = obj->data[bl];
                obj->data[bl] = obj->data[p];
                obj->data[p] = tmp;
                p = bl;
            } else {
                break;
            }
        }
    }
}

int kthLargestAdd(KthLargest* obj, int val) {
    if (obj->count < obj->capacity) {
        obj->data[obj->count] = val;
        heap_up(obj, obj->count);
        obj->count += 1;
    } else {
        if (val > obj->data[0]) {
            obj->data[0] = val;
            heap_down(obj, 0);
        }
    }
    return obj->count == obj->capacity ? obj->data[0] : -1;
}

KthLargest* kthLargestCreate(int k, int* nums, int numsSize) {
    int i;
    if (k < 1) {
        return NULL;
    }
    KthLargest* obj = malloc(sizeof(KthLargest));
    obj->data = malloc(sizeof(int) * k);
    obj->count = 0;
    obj->capacity = k;
    for (i = 0; i < numsSize; i++) {
        kthLargestAdd(obj, nums[i]);
    }
    return obj;
}

void kthLargestFree(KthLargest* obj) {
    free(obj->data);
    free(obj);
}

void printArr(int *a, int num) {
    for (int i = 0; i < num; i++) {
        printf("%4d", a[i]);
    }
    printf("\n");
}

int main(void) {
    int k = 3;
    int numsCount = 2;
    int param_1;
    int nums[] = {3, 2};
    KthLargest* obj = kthLargestCreate(k, nums, numsCount);
    
    printArr(obj->data, k);
    param_1 = kthLargestAdd(obj, 4);
    printf("---%d\n", param_1);
    printArr(obj->data, k);
    param_1 = kthLargestAdd(obj, 8);
    printf("---%d\n", param_1);
    printArr(obj->data, k);
    param_1 = kthLargestAdd(obj, 18);
    printf("---%d\n", param_1);
    printArr(obj->data, k);
    
    kthLargestFree(obj);

    return 0;
}