top-k算法的二叉堆数组实现(C++)（原创）

摘要：本文简要介绍了使用堆实现求序列中最大的K个元素的算法，并提供了C++源码

关键字：top-k，二叉堆，数组堆

本算法的具体思想为建立一个大小为K的最小堆（使用数组存储），然后将N个元素一次存于堆中，

每存放一次元素调整一次对，即替换对中的最小元素，并保持二叉堆为最小堆。当N个元素全部处理完时，

堆中的K个元素即为所求的最大K个元素。需要注意的是，最先向堆中插入K个元素时，直接将每个元素放入

堆数组的末尾，然后向上调整(move up)该元素，当向对中插入K+1~N个元素时，替换堆数组的第一个元素

（即堆中的最小元素），并向下(move down)调整该元素！该算法的时间复杂度为O(N * log2K)，特别适合

于数据不能全部装入内存的情形！

基于以上思想的top-k算法实现如下：

/*
 * top-k算法的数组堆实现
 * 使用数组存储的最小二叉堆，每次替换最小的元素
 * 时间复杂度nlogk
 * 不仅适用于数据全部在内存中的情形，
 * 还特别适用于数据不能全部装入内存，但内存中可以放入K个元素的情形
 * Compiler: VC++ 6.0
 */

#include <iostream>
#include <ctime>
using namespace std;

#define N 100000000
#define K 10

int array[N];
int heap[K];
void init_array(int n);
void print(int a[], int n);

int parent(int i);
void moveup(int i);
void movedown(int i);
void topk(int n, int k);

int main(int argc, char* *argv)
{
        srand(time(NULL));
        int n = 50;
        init_array(n);
        print(array, n);
        topk(n, K);
        print(heap, K);
        system("pause");
        return 0;
}

void init_array(int n)
{
        for(int i = 0; i < n; ++i)
                array[i] = rand() % 100;
}

void print(int a[], int n)
{
        printf("array:\n");
        for(int i = 0; i < n; ++i)
                cout<<a[i]<<" ";
        cout<<endl;
}

int parent(int i)
{
        if(i % 2 == 0) return (i / 2 -1);
        else return i / 2;
}

void moveup(int i)
{
        int sn = i;
        int prnt = parent(sn);
        while(prnt > 0)
        {
                if(heap[prnt] > heap[sn])
                {
                        swap(heap[prnt], heap[sn]);
                        sn = prnt;
                        prnt = parent(sn);
                }
                else break;
        }
}

void movedown(int i, int len)
{
        int j = i;
        while(2 * j + 1 < len)
        {
                int child = 2 * j + 1;
                if(child + 1 < len && heap[child + 1] < heap[child])
                        child++;
                //与较小的儿子节点交换
                if(heap[j] > heap[child])
                {
                        swap(heap[j], heap[child]);
                        j = child;
                }
                else break;
        }
}

void topk(int n, int k)
{
        int pos = 0;
        //将n个元素一次性放入堆中，每放一个的时间复杂度为logk
        for(int i = 0; i < n; ++i)
        {
                if(pos < k)
                {
                        heap[pos] = array[i];
                        moveup(pos++);
                }
                else
                {
                        if(heap[0] < array[i])
                        {
                                heap[0] = array[i];
                                movedown(0, k);
                        }
                }
        }
}