堆的簡單應用——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);

}