最大堆 最小堆 解决TOPK问题

堆：实质是一颗完全二叉树，最大堆的特点：父节点值均大于子节点；最小堆的父节点值均小于子节点；

 

一般使用连续内存存储堆内的值，因而可以根据当前节点的索引值推断子节点的索引值：

节点i的父节点为(i-1)/2；

节点j的左子结点：j * 2 + 1;

节点j的右子结点：j * 2 + 2;

 

以下代码实现了最大堆最小堆，当比较函数使用std::greater，得到最大堆，当比较函数使用std::less得到最小堆；

 

代码及测试用例如下：

 

//最大最小堆
//MaxMinHeap.h

#pragma once
#include <assert.h>

using namespace std;

template <typename T>
void mswap(T &a, T &b)
{
    T tmp = a;
    a = b;
    b = tmp;
}

template <typename T,typename Compare = std::less<T>>
class MaxMinHeap
{
public:
    int hSize ;    //堆空间
    int hCurNum;//堆内已占用空间
    T *data;

private:
    Compare comp;//比较函数
public:
    MaxMinHeap(int size)
    {
        hSize = size;
        assert(hSize>0);
        data = new T[hSize];
        hCurNum = 0;
    };
    ~MaxMinHeap(void)
    {
        if(data!=NULL)
            delete []data;
    };

    void headAdd(T num)
    {
        if (hCurNum==hSize)
        {
            if (comp(num,data[0]))//greater 大顶堆 保留最小的K个数；less 小顶堆 保留最大的K个数
                return;
            data[0]=num;
            HeapFixDown(0,hCurNum);
        }
        else
        {
            data[hCurNum++]=num;
            HeapFixUp(hCurNum-1);
        }
    };
    //最大堆排序后得到升序序列；最小堆排序后得到降序序列
    void sort()
    {
        for (int i=hCurNum-1; i >=1 ; --i)
        {
            mswap(data[i],data[0]);
            HeapFixDown(0,i);
        }
    }

    void GetHnum(T &n)//获取最大堆的最小值或者最小堆的最大值
    {
        n = data[0];
    };
    void HeapFixUp(int index)
    {
        assert (index < hCurNum);
        T tmp=data[index];
        int j = (index - 1)/2;//父节点
        while(j>=0 && index !=0)
        {
            if(comp(data[j],tmp))
                break;
            data[index]=data[j];
            index = j;
            j = (index - 1)/2;
        }
        data[index]=tmp;
    };

    //从节点index开始进行向下调整
    void HeapFixDown(int index, int n)
    {
        assert(index<hCurNum);
        assert(n<hCurNum);

        T tmp=data[index];
        int j = index*2+1;
        while(j<n)
        {
            if(j+1 < n && comp(data[j+1],data[j]))//大顶堆中左右孩子找最大的，小顶堆左右孩子找最小的
                ++j;
            if(comp(tmp,data[j]))
                break;
            data[index]=data[j];
            index = j;
            j = index*2+1;
        }
        data[index]=tmp;
    };
};

#include <functional>
#include <iostream>
#include "MaxMinHeap.h "

using namespace std;

int main(int argc ,  char ** argv)
{
    MaxMinHeap<float,greater<float>> test(20);

    for (int i = 0 ;i < 20; ++i)
    {
        test.headAdd(-i*2+38);
    }
    for (int i = 0 ; i < 20 ; ++i)
    {
        cout<<test.data[i]<<endl;
    }
    test.sort();
    for (int i = 0 ; i < 20 ; ++i)
    {
        cout<<test.data[i]<<" ";
    }
    cout<<endl;
    return 0;
}