當前位置: 首頁>>算法&結構>>正文


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);
                        }
                }
        }
}
本文由《純淨天空》出品。文章地址: https://vimsky.com/zh-tw/article/171.html,未經允許,請勿轉載。