堆的简单应用——TopK

1、海量数据top k问题

100亿个数中找出最大的前K个数，我们可以遍历K次找到，但是时间复杂度就很大为 O(KN);因此，我们可以用堆来实现，只需遍历一次，思路如下：

• 如果要找前K个最大的数，我们用小堆，每次用堆顶元素和遍历的数比，如果堆顶元素小，则让堆顶元素的值等于它，然后向下调整
• 如果要找前K个最小的数，我们用大堆，每次用堆顶元素和遍历的数比，如果堆顶元素大，则让堆顶元素的值等于它，然后向下调整

2、代码如下

这里我们暂时用有限的几个数模仿海量数据，来判断算法是否正确

TopK.h

#pragma once

//堆的操作见其他文章，这里只用到堆的结构体和，大堆小堆函数指针
#include "Heap.h"

//TopK
void TopK(Heap *hp, DataType * array, long size, int K);
//初始化堆
void HeapTopKInit(Heap *hp, Compare cmp, int K);
//交换元素
void SwopTopK(DataType *a, DataType *b);
//打印元素
void TopK_Print(Heap *hp, int K);
//向下调整算法
void AdjustDownTopK(Heap  *hp, DataType parent)
//测试
void TestTopK();

TopK.c

#include "TopK.h"

//求最小的3个元素
void TopK(Heap *hp, DataType * array, long size, int K)
{

    for (long i = 0; i < K; i++)
    {
        InsertHeap(hp, array[i]);
        AdjustDownTopK(hp, i);
    }
    for (long i = K; i < size; i++)
    {
        if (hp->cmp(array[0],hp->_array[i]))
        {
            SwopTopK(&hp->_array[0],&array[i]);
            AdjustDownTopK(hp, 0);
        }

    }
}

//向下调整
void AdjustDownTopK(Heap  *hp, DataType parent)
{
    int child = (parent<<1)+1;

    if (NULL == hp)
        return;
    while (child < hp->_size)
    {
        //找到孩子中较小的一个
        if ((child+1) < hp->_size && 
            hp->cmp(hp->_array[child+1], hp->_array[child]))
        {
            child += 1;
        }
        //如果双亲大于孩子，则交换
        if (hp->cmp(hp->_array[child], hp->_array[parent]))
        {
            Swop(&hp->_array[parent], &hp->_array[child]);
            parent = child;
            child = (parent << 1) + 1;
        }
        else
        {
            break;
        }
    }
}

void SwopTopK(DataType *a, DataType *b)
{
    DataType temp = *a;
    *a = *b;
    *b = temp;
}

void TopK_Print(Heap *hp, int K)
{
    int i = 0;
    for (; i < K; i++)
    {
        printf("%d ", hp->_array[i]);
    }
    printf("\n");
}

void HeapTopKInit(Heap *hp, Compare cmp, int K)
{
    if (NULL == hp)
        return;
    hp->_array = (DataType*)malloc(sizeof(DataType)* K);
    hp->_capacity = K;
    hp->_size = 0;
    hp->cmp = cmp;
}


void TestTopK()
{
    Heap hp;
    int array[] = { 53, 17, 78, 9, 45, 65, 87, 23, 31 };
    int K = 3;

    HeapTopKInit(&hp, Greater, K);
    TopK(&hp, array, sizeof(array)/sizeof(array[0]),K);
    TopK_Print(&hp, K);

}