TopK问题精讲-优快云博客

Top K问题：前K个最大的数

参考
1. 海量数据处理的 Top K算法(问题) 小顶堆实现
2. 《算法导论》读书笔记之第6章堆排序
3. 寻找最小的k个数
4. 编程之美

问题描述

输入N个整数， N很大，求出其中的前K个最大的数,K很小，一般在10以内。

由于输数据很大，但只求前K个最大的值，所以对所有输入数据的保存和排序是不可取的。

1 解法一：最小堆

利用最小堆的性质，每个非叶结点的数值，都小于等于其孩子节点的数值。利用含有K个节点的最小堆，保存K个目前的最大值。

每次接受新的数据输入时，先与根节点比较。若不大于根节点，则舍弃；否则用新数值替换根节点数值，并进行最小堆的调整。

1.1 最小堆复习

@算法导论

1.1.1 堆化数组

用数组存储堆，

若数组下标从1开始，则 $parent(i) = i/2, left(i) = i*2, right(i) = i*2+1$

若数组下标从0开始，则 $parent(i) = i/2, left(i) = i*2 + 1, right(i)=i*2+2$

最小堆的性质：

除了根结点以外，对于所有节点i，有 $A[parent(i)] \le A[i]$
含有n个元素的节点高度是ceil(lgn)
叶子节点的下标是n/2+1,…,n（下标从1开始）；n/2,…,n-1（下标从0开始）

以下讨论的都是下标从0开始的情况

1.1.2 维护堆

给定一个节点i，保证以i为根的子树满足堆性质，是自顶向下的过程

递归与非递归版的维护堆代码：

// 维护最小堆，递归, 调整以i为根的子树
void adjust_min_heap_recursive(int *a, int n, int i){
    int left = i * 2 + 1;  // 节点i的左孩子编号
    int right = left + 1; // 节点i的右孩子编号
    int smallest = i; // 默认最小是父节点
    if(left < n && a[left] < a[smallest]) // 找到i与其左右孩子中最小值的编号
        smallest = left;
    if(right < n && a[right] < a[smallest])
        smallest = right;
    if(smallest != i){ // 若最小值不是父节点
        swap(a[smallest], a[i]);
        adjust_min_heap_recursive(a, n, smallest); // 递归，继续自顶向下维护
    }
}

// 维护最小堆，非递归, 调整以i为根的子树
void adjust_min_heap(int *a, int n, int i){
    int left, right, smallest = i;
    while(1){
        left = 2 * i + 1;
        right = left + 1;
        if(left < n && a[left] < a[smallest])
            smallest = left;
        if(right < n && a[right] < a[smallest])
            smallest = right;
        if(smallest != i){
            swap(a[i], a[smallest]);
            i = smallest;
        } else {
            break;
        }
    }
}

1.1.3 建堆

建堆时，调用维护堆过程，是自底向上的过程，从数组的最后一个非叶结点(n/2-1)开始调整。

建堆代码：

// 建堆，堆化数组
void build_min_heap(int *a, int n){
    for(int i = n/2 - 1; i >= 0; --i){
        adjust_min_heap_recursive(a, n, i);
        // adjust_min_heap(a, n, i);
    }
}

1.1.4 堆排序（与Top K问题无关）

顺便复习一下堆排序

堆排序的流程是：

建最大堆，将输入数组a[1..N]转化为最小堆，使得最小值在堆顶，也就是a[0]；
每次不断将最大值提取出来，放置到数组的末尾，再将原数组末尾的值挪到堆顶，最大堆的size-1，并对剩下的对元素进行调整。

1.1.5 测试程序

#include <iostream>
#include <algorithm>
#include <cstdlib>
using namespace std;
// 数组下标从0开始

void print_array(int *a, int n){
    for(int i = 0; i < n; ++i){
        cout << a[i] << " ";
    }
    cout << endl;
}

// 维护最小堆，递归, 调整以i为根的子树，自顶向下
void adjust_min_heap_recursive(int *a, int n, int i){
    int left = i * 2 + 1;  // 节点i的左孩子编号
    int right = left + 1; // 节点i的右孩子编号
    int smallest = i; // 默认最小是父节点
    if(left < n && a[left] < a[smallest]) // 找到i与其左右孩子中最小值的编号
        smallest = left;
    if(right < n && a[right] < a[smallest])
        smallest = right;
    if(smallest != i){ // 若最小值不是父节点
        swap(a[smallest], a[i]);
        adjust_min_heap_recursive(a, n, smallest); // 递归，继续自顶向下维护
    }
}

// 建堆，堆化数组，自底向上
void build_min_heap(int *a, int n){
    for(int i = n/2 - 1; i >= 0; --i){
        adjust_min_heap_recursive(a, n, i);
        // adjust_min_heap(a, n, i);
    }
}

int main(){
    const int N = 10;
    int a[N];
    for(int i = 0; i < N; ++i){
        a[i] = rand() % 20;
    }
    cout << "堆化前数组：";
    print_array(a, N);

    build_min_heap(a, N);

    cout << "堆化后数组：";
    print_array(a, N);

    return 0;
}

1.2 Top K max 最小堆解法

时间复杂度O（nlogK）=O（n），空间复杂度O（K）=O（1）

前面第1节介绍的方法再提一遍。

利用含有K个节点的最小堆，保存K个目前的最大值。

每次接受新的数据输入时，先与根节点比较。若不大于根节点，则舍弃；否则用新数值替换当前的最小值，即根节点，并进行最小堆的调整。

#include <iostream>
#include <algorithm>
#include <fstream>
#include <cstdlib>
#include <cstdio>
using namespace std;

// 数组下标从0开始

void print_array(int *a, int n){
    for(int i = 0; i < n; ++i){
        cout << a[i] << " ";
    }
    cout << endl;
}


void write_rand_num(int n, int range)
{
    ofstream out("data.txt");
    if(out.is_open()){
        for(int i = 0; i < n; ++i){
            out << rand() % range << endl;
        }
        out.close();
    }
}

// 维护最小堆，递归, 调整以i为根的子树，自顶向下
// 时间复杂度O（logK），空间复杂度O(1)
void adjust_min_heap_recursive(int *a, const int n, int i){
    int left = i * 2 + 1;  // 节点i的左孩子编号
    int right = left + 1; // 节点i的右孩子编号
    int smallest = i; // 默认最小是父节点
    if(left < n && a[left] < a[smallest]) // 找到i与其左右孩子中最小值的编号
        smallest = left;
    if(right < n && a[right] < a[smallest])
        smallest = right;
    if(smallest != i){ // 若最小值不是父节点
        swap(a[smallest], a[i]);
        adjust_min_heap_recursive(a, n, smallest); // 递归，继续自顶向下维护
    }
}

// 建堆，堆化数组，自底向上
// 时间复杂度O（n）*O（logK） = O（nlogK），空间复杂度O(K)(建立的大小为K的堆化数组)
void build_min_heap(int *a, const int n){
    for(int i = n/2 - 1; i >= 0; --i){
        adjust_min_heap_recursive(a, n, i);
        // adjust_min_heap(a, n, i);
    }
}

int main(){

    int n = 100; // N很大
    // write_rand_num(n, 50); // 第一次运行，产生随机数据，写入data.txt

    // 第二次运行，从data.txt读数据
    const int K = 10; // K很小
    int *a = new int[K], temp;

    freopen("data.txt", "r", stdin);

    // 先读K个
    for(int i = 0; i < K; ++i){
        cin >> a[i];
    }
    // 堆化数组，建立最小堆
    build_min_heap(a, K);
    cout << "初始最小堆：";
    print_array(a, K);

    // 继续读剩余的n-K个
    for(int i = K; i < n; ++i){
        cin >> temp;
        if(temp > a[0]){
            a[0] = temp;
            adjust_min_heap_recursive(a, n, 0); // 向下调整
        }
    }

    cout << "最终最小堆：";
    print_array(a, K);

    fclose(stdin);
    return 0;
}

注意：必须分两次运行才可以，而且这样写实在太丑，目前猜测原因是两次写/读的流不同导致数据不同步。

另：是时候来个C/C++ IO专题了。

2 解法二：选择/交换排序

使用选择/交换排序，每轮对K个元素进行排序，流程如下：

存入前K个数，这K个数是当前最大的K个数。O（k）
对K个数，利用选择或交换排序找到K个元素中的最小值kmin。O（k）
继续遍历剩余n-k个数，设每次遍历到的新元素为x，

1) 若x>kmin, 则用x替换kmax，并回到第2步，找出k个元素中的最小值kmin’

2) 否则继续遍历下个元素。

时间复杂度O（n）*O（k） = O（nk）

使用快速排序，实现如下：

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <cstdlib>
using namespace std;

void print_array(int *a, int n){
    for(int i = 0; i < n; ++i){
        cout << a[i] << " ";
    }
    cout << endl;
}

// 划分，第一个数为基准
int partition(int *a, int from, int to)
{
    int i = from;
    int j = to + 1;
    int pivot = a[from];
    while(true){
        // 从前向后，找到第1个比基准大的数，i指向该数
        while(i < to && a[++i] < pivot);
        // 从后向前，找到第1个比基准小的数，j指向该数
        while(j > from && a[--j] > pivot);
        if(i >= j) break;
        swap(a[i], a[j]);
    }
    swap(a[from], a[j]);
    return j;
}

void quick_sort(int *a, int from, int to){
    if(from < to)
    {
        int p = partition(a, from, to);
        quick_sort(a, from, p - 1);
        quick_sort(a, p + 1, to);
    }
}

int main(){
    int n = 100;
    int K = 10;
    freopen("data.txt", "r", stdin);

    int *a = new int[K];
    int i, temp;

    for(i = 0; i < K; ++i){
        cin >> a[i];
    }
    quick_sort(a, 0, K-1);

    cout << "初始Top K max数组：";
    print_array(a, K);

    for(i = K; i < n; ++i){
        cin >> temp;
        if(temp > a[0]){
            a[0] = temp;
            quick_sort(a, 0, K-1);
        }
    }

    cout << "最终Top K max数组：";
    print_array(a, K);

    fclose(stdin);
    return 0;
}

寻找最大的K个数：Top K

编程之美2.5 寻找最大的K个数

元素数量有N个，选出其中最大的N个数.

当N很大的时候，比如说，N有100亿，怎么办？

首先想到的是：

1）可以将数据进行分成M块（M可以很大，利用多台机器，每台机器上还可以再分块），每块进行快速排序；

2）再从每块内取出最大的K个值，合成下一批待排序的 $M*K$ 个值，然后对 $M*K$ 个值继续排序

3）最终取排序后的最大的K个值。

因为无论怎样都要分块，所以又想到了M路归并：

1）将大数据数据文件划分成M块小数据数据文件。每块小数据进行快速排序，或者每块进行最小堆的Top K算法（最小堆大小为K，初始读入K个数据，之后再判断并读入剩余数据，并维护堆）。

2）内存中开M个队列（最小堆），每个队列从对应编号的小文件中读取数据，生成Top K个值。

3）在每个堆中取出1个形成一个中转站（最小堆），假设中转站大小为S（ $S>=K$ ），中转站从M个队列轮流取1个Top 1值，取值后执行最小堆的Top K算法。

4）最终取中转站中的Top K个值。

注意：下面讨论的解法比较适应非海量数据的环境，这些解法可作为上述方案中的内部排序（或选择）的算法。

3 解法三：排序

3.1 全排序

假设N不大，可以采用N个数排序的办法。

快速排序、堆排序均可，平均时间复杂度均为O（N*logN）
然后取出前K个，时间复杂度为O（K）

快速排序：

#include <iostream>
#include <algorithm>
#include <cstdlib>
using namespace std;

void print_array(int *a, int from, int to){
    for(int i = from; i <= to; ++i){
        cout << a[i] << " ";
    }
    cout << endl;
}

// 降序排列
int partition_random(int *a, int from, int to){
    int random_pos = from + rand() % (to - from + 1);
    swap(a[from], a[random_pos]);

    int i = from;
    int j = to + 1;
    int pivot = a[from];
    while(true){
        while(i < to && a[++i] > pivot);
        while(j > from && a[--j] < pivot);
        if(i >= j)
            break;
        swap(a[i], a[j]);
    }
    swap(a[from], a[j]);
    return j;
}

void quick_sort(int *a, int from, int to){
    if(from < to)
    {
        int p = partition_random(a, from, to);
        quick_sort(a, from, p-1);
        quick_sort(a, p+1, to);
    }
}

int main(){
    int n = 1000;
    int K = 10;
    int *a = new int[n];
    for(int i = 0; i < n; ++i){
        a[i] = rand() % n;
    }

    quick_sort(a, 0, n-1);
    print_array(a, 0, K-1);
    delete[] a;
    return 0;
}

堆排序，N个数全部排序，最小堆：

#include <iostream>
#include <algorithm>
#include <cstdlib>
using namespace std;

void print_array(int *a, int from, int to){
    for(int i = from; i <= to; ++i){
        cout << a[i] <<" ";
    }
    cout << endl;
}
void adjust_min_heap(int *a, int i, int n){
    int left = 2 * i + 1;
    int right = left + 1;
    int smallest = i;
    if(left < n && a[left] < a[smallest]) smallest = left;
    if(right < n && a[right] < a[smallest]) smallest = right;
    if(smallest != i){
        swap(a[i], a[smallest]);
        adjust_min_heap(a, smallest, n);
    }
}


void build_min_heap(int *a, int n){
    for(int i = n/2 - 1; i >= 0; --i){
        adjust_min_heap(a, i, n);
    }
}

void heap_sort(int *a, int n){
    build_min_heap(a, n);
    for(int i = n - 1; i >= 0; --i){
        swap(a[0], a[i]);
        adjust_min_heap(a, 0, i);
    }
}

int main()
{
    int n = 1000;
    int K = 10;
    int *a = new int[n];
    for(int i = 0; i < n; ++i){
        a[i] = rand() % n;
    }
    heap_sort(a, n);
    print_array(a, 0, K-1);
    return 0;
}

总时间复杂度O（N*logN），空间复杂度O（N）。

3.2 K个数排序

上面的方法，时间复杂度为O（NlogN），与K无关。若K < logN，则可进一步避免做后面N-K个数的排序。

选择排序或交换排序均可，把N个数中的前K大个数排序，时间复杂度为O（N*K）

选择排序：选择最大的K数

#include <iostream>
#include <algorithm>
#include <cstdlib>
using namespace std;

void print_array(int *a, int from, int to){
    for(int i = from; i <= to; ++i){
        cout << a[i] <<" ";
    }
    cout << endl;
}

// 选择前K大的选择排序
void select_sort(int *a, int n, int K){
    int index;
    for(int i = 0; i < K; ++i){
        index = i;
        //
        for(int j = i+1; j < n; ++j){
            if(a[j] > a[index]) index = j;
        }
        if(index != i) swap(a[i], a[index]);
    }
}

int main()
{
    int n = 1000;
    int K = 10;
    int *a = new int[n];
    for(int i = 0; i < n; ++i){
        a[i] = rand() % n;
    }
    select_sort(a, n, K);
    print_array(a, 0, K-1);
    return 0;
}

冒泡排序，选择最大的K个数：

#include <iostream>
#include <algorithm>
#include <cstdlib>
using namespace std;

void print_array(int *a, int from, int to){
    for(int i = from; i <= to; ++i){
        cout << a[i] <<" ";
    }
    cout << endl;
}

void bubble_sort(int *a, int n){
    for(int i = 0; i < n - 1; ++i){ // 每趟排序，最小值“沉底”，所以共需n-1趟，
        for(int j = 0; j < n - i - 1; ++j){ // 每趟排序中，从头开始，两两交换，后面有i个元素已“沉底”
            if(a[j] < a[j+1]) swap(a[j], a[j+1]);
        }
    }
}

int main()
{
    int n = 1000;
    int K = 10;
    int *a = new int[n];
    for(int i = 0; i < n; ++i){
        a[i] = rand() % n;
    }
    bubble_sort(a, n);
    print_array(a, 0, K-1);
    return 0;
}

4 解法四：快速选择

下面继续优化，避免对前K个数进行排序，进一步提升性能。

修改快速排序算法，得到快速选择算法，快速选择算法流程如下：

假设N个数存在数组S中，从S中随机找出一个元素X，把数组分成两部 $S_a$ 和 $S_b$

$S_a$ 中的元素大于等于X， $S_b$ 中的元素小于X。

这时有两种可能：

$|S_a|<K$ ，则 $S_a$ 中所有的数和 $S_b$ 中最大的 $K-S_a$ 个元素就是S中最大的K个数.
$|S_a|\ge K$ ，则要返回 $S_a$ 中最大的K个元素。
递归，不断把问题分解成更小的问题，平均时间复杂度O（N*logK）

#include <iostream>
#include <algorithm>
#include <cstdlib>
using namespace std;


void print_array(int *a, int from, int to){
    for(int i = from; i <= to; ++i){
        cout << a[i] <<" ";
    }
    cout << endl;
}

int partition(int *a, int from, int to){
    int rand_pos = from + rand() % (to - from + 1);
    swap(a[from], a[rand_pos]);

    int i = from;
    int j = to + 1;
    int pivot = a[from];
    while(1){
        while(i < to && a[++i] > pivot);
        while(j > from && a[--j] < pivot);
        if(i >= j) break;
        swap(a[i], a[j]);
    }
    swap(a[from], a[j]);
    return j;
}

void quick_select(int *a, int k, int from, int to){
    if(from < to){
        int p = partition(a, from, to);
        int sa = p - from;
        if(sa > k)
            quick_select(a, k, from, p - 1);
        else if(sa < k)
            quick_select(a, k-sa-1, p + 1, to);
        else
            return;
    }
}

int main()
{
    int n = 11;
    int k = 5;
    int a[] = {1,2,3,4,5,6,7,8,9, 9,10};

    quick_select(a, k, 0, n - 1);
    print_array(a, 0, k-1);
    return 0;
}