top-k 算法淺析

簡介

top-k算法，其實就是尋找最大的k個數（具體詳見《編程之美》第2.5節「尋找最大的k個數」）。好比一個數組：1，2，5，9，4，3，7 須要尋找最大的2個數，那麼就是9和7。最先以前我接觸到topk算法的時候，以爲解決思路就是排序，排完序以後，取前k個數就能夠了。可是這種思路雖然簡單，可是效率是不好的。由於題目只要求最大的k個數，並不須要k個數有序，也不須要後n-k個數有序。

解決方法

我用的是解法四，用一個容量爲k的最小堆來儲存最大的k個數。最小堆的堆頂元素就是最大k個數中最小的一個。每次新考慮一個數x，若是x比堆頂的元素y小，則不須要改變原來的堆，由於這個元素比最大的k個數小。若是x比堆頂的元素大，那麼用x替換堆頂的元素y。在x替換堆頂元素y後，x可能破壞最小堆的結構（每一個結點都比它的父親結點大），須要更新堆來維持堆的性質。

代碼實現（C語言）

#include <stdio.h>
#include <stdlib.h>

// 定義取最大N個數
#define TOP_K 6
int heap[6];

// 交換數據
void swap(int *a, int *b)
{
    int temp = *b;
    *b = *a;
    *a = temp;
}

// 調整最小堆
void min_heapify(int arr[], int start, int end)
{
    int dad = start;
    int son = dad * 2 + 1;
    while (son <= end)
    {
        if (son + 1 <= end && arr[son] > arr[son + 1])
            son++;
        if (arr[dad] < arr[son])
            return;
        else
        {
            swap(&arr[dad], &arr[son]);
            dad = son;
            son = dad * 2 + 1;
        }
    }
}

// 創建最小堆
void buid_heap(int heap[])
{
    int i;
    for (i = TOP_K / 2; i >= 0; i--)
    {
        min_heapify(heap, i, TOP_K - 1);
    }
}

// 8,8,8,9,9,9
int main()
{
    int arr[] = {3, 5, 3, 0, 8, 6, 1, 5, 8, 6, 2, 4, 9, 4, 7, 0, 1, 8, 9, 7, 3, 1, 2, 5, 9, 7, 4, 0, 2, 6};
    int len = (int)sizeof(arr) / sizeof(*arr);
    int i;

    // 堆賦值
    for (i = 0; i < TOP_K; i++)
    {
        heap[i] = arr[i];
    }

    buid_heap(heap); // 創建最小堆

    // 循環遍歷整個數組
    for (i = TOP_K + 1; i <= len; i++)
    {
        if (arr[i] > heap[0]) // 只有大於根節點才處理
        {
            heap[0] = arr[i];
            min_heapify(heap, 0, TOP_K - 1); // 向下調整堆
        }
    }

    // 打印最大key個數
    for (i = 0; i < TOP_K; i++)
    {
        printf("%d ", heap[i]);
    }
}