第K大/Top K及其简单实现

转载请注明出处：http://blog.csdn.net/u012469987/。

---

见网上第K大多数只给思路，没给实现，我就来填坑了。
Top K 和第K大基本等价，以下我们以第K大为例且假设第K大一定存在，Top K 可以在第k大基础上稍微改动获得。
本文介绍6种方法，只考虑实现功能，不做异常判断，面试的话快排和最小堆的方法比较不错，测试提交的话可以去Leetcode，或者直接拿最下面的数据生成代码去对拍跑。

快排的思想 近似O(n)

调用降序快排的partition函数，设区间为[low,high]，返回index，则index左边都是大于data[index]的。
1. 若index及index左边数字有k个则data[index]就是第k大，index及其左边元素为Top K元素
2. 左边数字大于k个则继续在[low,index]里找
3. 左边数字小于k个则去右边[index+1,high]找 k - 左边数字个数

#include <cstdio>#include <iostream>using namespace std;const int maxn = 1e5 + 5;//改为 data[high] >= key 和 data[low] <= key 则为第k小int part(int *data, int low, int high) {    int key = data[low];    while (low < high) {        while (low < high && data[high] <= key) high--;        data[low] = data[high];        while (low < high && data[low] >= key) low++;        data[high] = data[low] ;    }    data[low] = key;    return low;}int k_th(int *data, int k, int low, int high) {    int pos = part(data, low, high);    int cnt = pos - low + 1;  //[low,pos]元素个数    if (cnt == k) return data[pos];    else if (cnt < k) return k_th(data, k - cnt, pos + 1, high);    else return k_th(data, k, low, pos);}int k_th(int *data, int n, int k) {    if(k < 1 || k > n) return -1;    return k_th(data, k, 0, n - 1); //闭区间    //遍历data[0,k)即可获得top K，且是有序的}int main() {    // int data[] = {1, 5, 6, 7, 3, 2, 10, 9, 0, 231, 3214, 61};    // int n = sizeof(data) / sizeof(int);    // int k = 2;    // cout << k_th(data, n, k) << endl;    // freopen("in.txt","r",stdin);    // freopen("out.txt","w",stdout);    int n, k, data[maxn];    std::ios::sync_with_stdio(false);    while (cin >> n >> k) {        for (int i = 0; i < n; ++i) {            cin >> data[i];        }        cout << k_th(data, n, k) << endl;    }    return 0;}

小根堆 O(nlogk)

维护一个k个元素的小根堆，保持堆顶为第k大，扫一遍数据，若堆里个数小于k则插入，否则看新的数和堆顶数大小关系：
1. 若新来的数小于等于堆顶，即新元素比Top K里最小的还小，则新来的数显然不可能是前k大
2. 若新来的数大于堆顶，则删掉堆顶，将新数字放到堆里且调整堆来保持堆的属性

由于实现堆代码量较多，我们可以用C++的优先队列、set等代替手工堆偷跑，当然这里也提供了手动实现版。

#include <cstdio>#include <vector>#include <queue>#include <iostream>using namespace std;const int maxn = 1e5 + 5;//维持一个k大小的最小堆，根据新元素和堆顶大小决定要不要加入堆且删堆顶// O(nlogk)int biggest_k_th(int *data, int n, int k) {    priority_queue<int, vector<int>, greater<int> >q;   //小根堆    while (!q.empty()) q.pop();    for (int i = 0; i < n; ++i) {        if (q.size() < k) {            q.push(data[i]);        } else if (data[i] > q.top()) {            q.pop();            q.push(data[i]);        }    }    //取k次q.top()且pop()k次即为有序的前K大    return q.top();}int smallest_k_th(int *data, int n, int k) {    priority_queue<int>q;   //大根堆    while (!q.empty()) q.pop();    for (int i = 0; i < n; ++i) {        if (q.size() < k) {            q.push(data[i]);        } else if (data[i] < q.top()) {            q.pop();            q.push(data[i]);        }    }    return q.top();}int main() {    // freopen("in.txt","r",stdin);    // freopen("out.txt","w",stdout);    std::ios::sync_with_stdio(false);    int n, k, data[maxn];    while (cin >> n >> k) {        for (int i = 0; i < n; ++i) {            cin >> data[i];        }        cout << biggest_k_th(data, n, k) << endl;    }    return 0;}

手动实现版

#include <iostream>#include <cstdio>#include <cstring>using namespace std;const int maxn = 1e5 + 5;const int maxK = 1e5 + 5;int heapCnt = 0;int heap[maxK];void adjust(int *heap, int begin, int end) {    //[begin,end)    int cur = begin;    int son = 2 * cur + 1;    while (son < end) {        if (son + 1 < end && heap[son] > heap[son + 1]) son++;        if (heap[cur] < heap[son]) return;        swap(heap[son], heap[cur]);        cur = son;        son = 2 * cur + 1;    }}void buildHeap(int *heap, int k) {  //[data,data+k) 开区间    for (int i = k / 2; i >= 0;  --i) {        adjust(heap, i, k);    }}int k_th(int *data, int n, int k) {    heapCnt = 0;    for (int i = 0; i < n; ++i) {        if (heapCnt < k) {            heap[heapCnt++] = data[i];            if (heapCnt == k) {                buildHeap(heap, k); //data[0,k)共k个            }        } else {            if (data[i] > heap[0]) {                heap[0] = data[i];                adjust(heap, 0, heapCnt);            }        }    }    return heap[0];}int main() {    // freopen("in.txt", "r", stdin);    // freopen("out.txt", "w", stdout);    int n, k, data[maxn];    std::ios::sync_with_stdio(false);    while (cin >> n >> k) {        for (int i = 0; i < n; ++i) {            cin >> data[i];        }        cout << k_th(data, n, k) << endl;    }    return 0;}

计数排序 O(n)

按照计数排序思想给数据的值计数，再从数据的最大值往最小值遍历，则总次数大于等于k的那个数为第k大，见代码一目了然。
优点：速度快且不用库也代码量少，妥妥的O(n)
缺点：只适用于数值不大的情况，当然你用hashmap这类库计数的话就没这问题了。

#include <iostream>#include <cstdio>#include <cstring>using namespace std;const int maxn = 1e5 + 5;const int maxVal = 1e5 + 5; //O(n) 适用于数据值不大的情况int k_th(int *data, int n, int k) {    int mmin = data[0], mmax = data[0];    int times[maxVal];    memset(times,0,sizeof(times));    for (int i = 0; i < n; ++i) {        mmin = min(mmin, data[i]);        mmax = max(mmax, data[i]);        times[data[i]]++;    }    int cnt = 0;    for (int i = mmax; i >= mmin; --i) {        cnt += times[i];        if (cnt >= k) { // >= 是因为第k大的数可能有若干个            return i;        }        //反过来遍历则为第k小        //每次输出times[i]次i则为有序前k大    }    return -1;}int main() {    // freopen("in.txt","r",stdin);    // freopen("out.txt","w",stdout);    int n, k, data[maxn];    std::ios::sync_with_stdio(false);    while (cin >> n >> k) {        for (int i = 0; i < n; ++i) {            cin >> data[i];        }        cout<< k_th(data, n, k) <<endl;    }    return 0;}

二分 O(nlogn)

二分答案值区间[l,r]，最开始l=所有数的最小值，r=最大值，假设当前值是mid，如果所有数据中大于等于mid的数字至少k个，说明当前数值可能是答案（若mid存在的情况则将区间调为[mid,r]，mid不存在的话就改为[mid+1,r]），否则mid偏大，在[l,mid-1]里查找；二分不会的可见这篇文章。

#include <iostream>#include <cstdio>#include <cstring>using namespace std;const int maxn = 1e5 + 5;const int maxVal = 1e5 + 5;bool ok(int *data, int n, int k, int mid) {    int cnt = 0;    for (int i = 0; i < n; ++i) {        if (data[i] >= mid) cnt++;    }    return cnt >= k;}int k_th(int *data, int n, int k) {    int mmin = data[0], mmax = data[0];    bool vis[maxVal];    memset(vis, false, sizeof(vis));    for (int i = 0; i < n; ++i) {        mmin = min(mmin, data[i]);        mmax = max(mmax, data[i]);        vis[data[i]] = true;    }    int l = mmin, r = mmax;    while (l < r) {        int mid = (l + r + 1) / 2;        if (ok(data, n, k, mid)) {            if (!vis[mid]) l = mid + 1;            else l = mid;        } else {            r = mid - 1;        }    }    return l;}int main() {    // freopen("in.txt", "r", stdin);    // freopen("out.txt", "w", stdout);    int n, k, data[maxn];    std::ios::sync_with_stdio(false);    while (cin >> n >> k) {        for (int i = 0; i < n; ++i) {            cin >> data[i];        }        cout << k_th(data, n, k) << endl;    }    return 0;}

暴力式选择/冒泡排序 O(kn)

特慢做法：排序k个，每次遍历n个元素，O(k*n)

#include <iostream>#include <cstdio>#include <cstring>using namespace std;const int maxn = 1e5 + 5;int k_th(int *data, int n, int k) {    for (int i = 0; i < k; ++i) {        for (int j = 0; j < n - i - 1; ++j) {            if (data[j] > data[j + 1]) {                swap(data[j], data[j + 1]);            }        }    }    return data[n-k];}int main() {    // freopen("in.txt", "r", stdin);    // freopen("out.txt", "w", stdout);    int n, k, data[maxn];    std::ios::sync_with_stdio(false);    while (cin >> n >> k) {        for (int i = 0; i < n; ++i) {            cin >> data[i];        }        cout << k_th(data, n, k) << endl;    }    return 0;}

真暴力排序O(nlogn)

排完取 data[k]，这么暴力就不说了。

数据生成代码

生成10组数据，每组一个n（范围：[a_n,b_n]），然后n个数 [a,b]。

#include <cstdio>#include <cmath>#include <cstdlib>using namespace std;int rand_ab(int a, int b) { //[a,b]    return a + rand() % (b + 1 - a);}void make(){        int a_n = 10000, b_n = 100000;    int a = 1, b = 10000;    for (int i = 0; i < 10; ++i) {        int n = rand_ab(a_n, b_n);        printf("%d ", n);        int a_k = 1, b_k = n;        printf("%d\n", rand_ab(a_k,b_k));        printf("%d", rand_ab(a, b));        for (int i = 1; i < n; ++i) {            printf(" %d", rand_ab(a, b));        }        printf("\n");    }}int main() {    // freopen("out.txt","w",stdout);    make();    return 0;}