在 1 亿个数中取出最大的 1000 个数？

起因

上个月的一场面试中，面试官问了我这个问题，当时没能很好的回答出来。后来认真的思考了下这个问题，用此文记录下来。

排序算法

符合第一直觉的回答应该是快速排序，在系统的学过排序算法后，我们都知道快排是一种很快的排序方法。我们先对数据进行排序，然后再取出其中前 1000 个数。不过你要是真的在面试中只那么回答的话，那么基本也无缘后面的面试了。这种思路在输入的数据规模很小时，不失为一种便捷的方法，但是当数据规模扩大时（题目中的 1 亿），其对运行的内存要求也将增大。在快排中，其时间复杂度为 O(n log n)，要同时对 1 亿整数排序，需要将其读入内存中，以每个 int 类型占 4 字节为例，1 亿整数则需要占用 400MB 的内存。虽然现代的计算机可用内存远远大于400MB，但当数据规模再次扩大到 10 亿甚至是 100 亿时，其对内存的要求则上升到了 4GB 和 40GB，算法的可用性 (availability) 和扩展性 (scalability) 将受到严重的影响。同时，在这种情况下还需要做全排序的话，就只能借助外部排序、位图排序、基数排序、桶排序 等算法来解决内存不足的问题了。同时，题目中只要求取 top 1000 的数据，对所有元素排序很明显是不必要的。

我写一个简单的演示程序，代码如下：

#include <iostream>  
#include <algorithm>  
#include <random>  
#include <chrono>  
  
using namespace std;  
  
#define MAX INT32_MAX  
#define MIN INT32_MIN  
#define SIZE 100000000 // 100 million  
#define K 1000  
  
int main() {   
    auto start = chrono::high_resolution_clock::now();  

	  // Create a random number generator 
    random_device rd;  
    mt19937 gen(rd());  
    uniform_int_distribution<> dis(MIN, MAX);  
  
    // Create an array of random numbers  
    int *arr = new int[SIZE];  
    for (int i = 0; i < SIZE; i++) {  
        arr[i] = dis(gen);  
    }  
  
    auto end = chrono::high_resolution_clock::now();  
    chrono::duration<double> diff = end - start;  
    cout << "Time taken to create array: " << diff.count() << " s" << endl;  
  

    start = chrono::high_resolution_clock::now();  
	// sorting via std::sort() - optimized quicksort
    sort(arr, arr + SIZE);  
  
    end = chrono::high_resolution_clock::now();  
    diff = end - start;  
    // the arr[0]~arr[k-1] is the top k elements in arr.
    cout << "Time taken to get top K elements: " << diff.count() << " s" << endl;
  
    free(arr);  
    return 0;  
}

# CPU: i7-8700k
# Memory: 16 GB
# OS: Ubuntu 22.10 Kinetic
Time taken to create array: 2.13481 s
Time taken to get top K elements: 22.3849 s

局部淘汰法

为了避免上面的算法中的不必要排序，我们可以选取出待排序数组中的前 k (= 1000) 个元素构建一个最小堆。然后从 k + 1 开始与容器中的最小元素 m 相比，如果大于 m，则替换 m 并重新调整最小堆。

其算法如下：

选取前K个元素，构建一个最小堆。
从第 K+1 个元素开始，与堆顶元素比较，如果比堆顶元素大，则替换堆顶元素，并调整堆。
重复步骤 2，直到数组遍历完毕。
堆中的元素即为前K个最大元素。

其代码如下：

#include <iostream>  
#include <algorithm>  
#include <random>  
#include <chrono>  
  
using namespace std;  
  
#define MAX INT32_MAX  
#define MIN INT32_MIN  
#define SIZE 100000000 // 100 million  
#define K 1000  
  
int main() {  
    // Create a random number generator  
    auto start = chrono::high_resolution_clock::now();  
  
    random_device rd;  
    mt19937 gen(rd());  
    uniform_int_distribution<> dis(MIN, MAX);  
  
    // Create an array of random numbers  
    int *arr = new int[SIZE];  
    for (int i = 0; i < SIZE; i++) {  
        arr[i] = dis(gen);  
    }  
  
    auto end = chrono::high_resolution_clock::now();  
    chrono::duration<double> diff = end - start;  
    cout << "Time taken to create array: " << diff.count() << " s" << endl;  
  
    // clock start  
    start = chrono::high_resolution_clock::now();  

	// generate a min-heap with first k elements
    int *heap = new int[K];  
    for (int i = 0; i < K; i++) {  
        heap[i] = arr[i];  
    }
    make_heap(heap, heap + K, greater<>());  
    for (int i = K; i < SIZE; i++) {  
        if (arr[i] > heap[0]) {  
            pop_heap(heap, heap + K, greater<>());  
            heap[K - 1] = arr[i];  
            push_heap(heap, heap + K, greater<>());  
        }  
    }  
  
    // clock end  
    end = chrono::high_resolution_clock::now();  
    diff = end - start;  
    cout << "Time taken to get top K elements: " << diff.count() << " s" << endl;  
  
    free(arr);  
    free(heap);  
    return 0;  
}

Time taken to create array: 2.18854 s
Time taken to get top K elements: 0.178133 s

C++ STL `set`

我们知道，标准库中的 set 底层实现是一棵红黑树（平衡二叉树），在这里我想到了是不是可以借助 set 容器来取 top k 元素呢。如果允许重复元素的话，则使用 multiset 来代替 set.

代码如下：

#include <iostream>  
#include <random>  
#include <chrono>  
#include <set>  
  
using namespace std;  
  
#define MAX INT32_MAX  
#define MIN INT32_MIN  
#define SIZE 100000000 // 100 million  
#define K 1000  
  
int main() {  
    // Create a random number generator  
    auto start = chrono::high_resolution_clock::now();  
  
    random_device rd;  
    mt19937 gen(rd());  
    uniform_int_distribution<> dis(MIN, MAX);  
  
    // Create an array of random numbers  
    int *arr = new int[SIZE];  
    for (int i = 0; i < SIZE; i++) {  
        arr[i] = dis(gen);  
    }  
  
    auto end = chrono::high_resolution_clock::now();  
    chrono::duration<double> diff = end - start;  
    cout << "Time taken to create array: " << diff.count() << " s" << endl;  
  
    // clock start  
    start = chrono::high_resolution_clock::now();  
  
    // Create a set of K the largest numbers  
    set<int> s;  
    for (int i = 0; i < SIZE; i++) {  
        if (s.size() < K) {  
            s.insert(arr[i]);  
        } else {  
            if (arr[i] > *s.begin()) {  
                s.erase(s.begin());  
                s.insert(arr[i]);  
            }  
        }  
    }  
  
    // clock end  
    end = chrono::high_resolution_clock::now();  
    diff = end - start;  
    cout << "Time taken to get top K elements: " << diff.count() << " s" << endl;  
  
    // Print the set in reverse order  
//    for (auto it = s.rbegin(); it != s.rend(); it++) {  
//        cout << *it << endl;  
//    }
  
    free(arr);  
    return 0;  
}

Time taken to create array: 2.16181 s
Time taken to get top K elements: 1.68677 s

总结

通过上面的实验，我们可以看到，使用局部淘汰的思想和最小堆的数据结构来取 top k 元素的效率是最高的，其次是使用 set 容器，最后是使用排序算法。

参考

1 亿个数中找出最大的100个数 (top K问题) - CSDN

最后修改于 2023-02-03

Pintaste

起因

排序算法

局部淘汰法

C++ STL `set`

总结

参考

Pintaste

Pintaste

起因

排序算法

局部淘汰法

C++ STL set

总结

参考

Pintaste

C++ STL `set`