数据结构之堆的应用—TopK问题

Posted 2021-09-19 小赵小赵福星高照～

TopK问题

Top-K问题：即求数据结合中前K个最大的元素或者最小的元素，一般情况下数据量都比较大。

这里我们讲解三种方法：

文章目录

• TopK问题
• • 方法一：排序
  • 方法二：将N个数建堆，取出前K个
  • 方法三：建一个K个数的堆(最优)

方法一：排序

思路：

将所有数据进行排序，取前K个元素即可

void swap(int *p1,int *p2)
{
    int temp=*p1;
    *p1=*p2;
    *p2=temp;
}
void AdjustDown(int *a,int n,int parent)
{
    int child=parent*2+1;
    while(child<n)
    {
        if(child+1<n && a[child+1]<a[child])//右孩子存在且右孩子小于或大于左孩子(建小堆找小的孩子，大堆找大的孩子)
        {
            child=child+1;
        }
        if(a[parent]>a[child])
        {
            swap(&a[parent],&a[child]);
            parent=child;
            child=parent*2+1;
        }
        else
        {
            break;
        }
    }
}
void HeapSort(int *a,int n)
{
    //建堆
    for(int i=(n-1-1)/2;i>=0;i--)
    {
        AdjustDown(a,n,i);
    }
        //排序
    int end = n-1;
    while(end>0)
    {
        swap(&a[0],&a[end]);
        end--;
        AdjustDown(a,end,0);
    }
}
void PrintTopK(int *a,int n,int k)
{
    HeapSort(a,n);
    for(int i=0;i<k;i++)
    {
        printf("%d ",a[i]);
    }
    printf("\\n");
}

时间复杂度O(N*logN)

方法二：将N个数建堆，取出前K个

思路：

建好堆，(找最大的前K个建大堆，找最小的前K个建小堆)然后将堆顶的数与最后一个数交换位置，向下调整一次找次大/小的数(注意向下调整时不要包含交换后的最后一个位置)，重复这样的操作，直到找完K个数

void swap(int *p1,int *p2)
{
    int temp=*p1;
    *p1=*p2;
    *p2=temp;
}
void AdjustDown(int *a,int n,int parent)
{
    int child=parent*2+1;
    while(child<n)
    {
        if(child+1<n && a[child+1]<a[child])//右孩子存在且右孩子小于或大于左孩子(建小堆找小的孩子，大堆找大的孩子)
        {
            child=child+1;
        }
        if(a[parent]>a[child])
        {
            swap(&a[parent],&a[child]);
            parent=child;
            child=parent*2+1;
        }
        else
        {
            break;
        }
    }
}
void PrintTopK(int *a,int n,int k)
{
    //建堆
    for(int i=(n-1-1)/2;i>=0;i--)
    {
        AdjustDown(a,n,i);
    }
    int end=n-1;
    for(int i=0;i<k;i++)
    {
        printf("%d ",a[0]);
        swap(&a[0],&a[end]);//交换最值和最后一个位置的数
        end--;//不包含最后一个位置
        AdjustDown(a,end,0);//调整
    }
}

时间复杂度O(N+K*logN)

方法三：建一个K个数的堆(最优)

TOP-K问题：即求数据结合中前K个最大的元素或者最小的元素，一般情况下数据量都比较大。

对于Top-K问题，能想到的最简单直接的方式就是排序，但是：如果数据量非常大，排序就不太可取了(可能数据都不能一下子全部加载到内存中)。最佳的方式就是用堆来解决，基本思路如下：

• 用数据集合中前K个元素来建堆
  前k个最大的元素，则建小堆
  前k个最小的元素，则建大堆

• 用剩余的N-K个元素依次与堆顶元素来比较，不满足则替换堆顶元素

将剩余N-K个元素依次与堆顶元素比完之后，堆中剩余的K个元素就是所求的前K个最小或者最大的元素。

思路：

首先我们建一个K个数的堆，需要前K个最大的数时，建小堆。需要前K个最小的数时，建大堆。然后将n-k个数依次与堆顶比较，如果比堆顶小(需要前K小个数)/大(需要前K大个数)，则进K个数的堆，然后调整找次大/小的数

注意：

需要前K个最大的数时，建小堆

需要前K个最小的数时，建大堆

#include<stdio.h>
void swap(int* p1, int* p2)
{
    int temp = *p1;
    *p1 = *p2;
    *p2 = temp;
}
void AdjustDown(int* a, int n, int parent)
{
    int child = parent * 2 + 1;
    while (child < n)
    {
        if (child + 1 < n && a[child + 1] < a[child])//右孩子存在且右孩子小于或大于左孩子(建小堆找小的孩子，大堆找大的孩子)
        {
            child = child + 1;
        }
        if (a[parent] > a[child])
        {
            swap(&a[parent], &a[child]);
            parent = child;
            child = parent * 2 + 1;
        }
        else
        {
            break;
        }
    }
}
void PrintTopK(int* a, int n, int k)
{
    //建k个数的堆
    for (int i = (k - 1 - 1) / 2; i >= 0; i--)//O(k)
    {
        AdjustDown(a, k, i);//取最大的需要建小堆，取最小的需要建大堆
    }
    for (int i = k; i < n; i++)//O(n*logk)
    {
        if (a[0] < a[i])//比堆顶大的就入堆
        {
            swap(&a[0], &a[i]);
            AdjustDown(a, k, 0);
        }
    }
    for (int i = 0; i < k; i++)
    {
        printf("%d ", a[i]);
    }
}

时间复杂度O(K+N*logK)，当我们要取前K个最大的数时，建一个K个数的小堆，剩下的数比堆顶大，就替换堆顶的数据进堆当最大的前K个数，一些数没进堆，那么说明其它数在堆里面，这些数肯定比topK要小，当最大的K数中某一个来了以后，一定比堆顶的数据大，一定能进堆