如何在3s内对10亿(1G)正整数进行排序

很早之前我曾经写过一篇博客《基数排序的性能优化》，详细介绍了如何对基数排序进行优化，使其速度能达到快排的5倍左右。在博客的最后也提到，可以通过修改基数实现对任意int型正整数的排序。在此，我们继续尝试通过并行手段优化基数排序，速度最快可达系统排序的20倍左右，最终可以在3s左右完成10亿（也即1B）int型正整数的排序。

1.非并行实现

首先我们假定数据的范围为0~2^31-1范围内的正整数，如果我们选择基数为2048，则循环3次即可覆盖整个范围；如果我们选择基数为65536，则只需要2次循环即可覆盖整个范围。当数据量很大时，我们可以选择65536作为基数。我们先给出非并行代码。
基数为2048：

void radix_sort_2048(int *arr, int n) {
    int *backup = new int[n];
    int mod = 0x7ff;

    for (int i = 0; i < 3; i++) {
        int bucket[2048] = {0};
        int move = i * 11;
        for (int j = 0; j < n; j++) {
            bucket[(arr[j] >> move) & mod]++;
        }

        for (int j = 1; j < 2048; j++) {
            bucket[j] += bucket[j - 1];
        }

        for (int j = n - 1; j >= 0; j--) {
            backup[--bucket[(arr[j] >> move) & mod]] = arr[j];
        }

        swap(backup, arr);
    }

    swap(backup, arr);//需要再次交换，避免删除原始数组

    memcpy(arr, backup, n*sizeof(int)); //3次循环，需要再多一次拷贝

    delete[] backup;
}

基数为65536：

void radix_sort_65536(int *arr, int n) {
    int *backup = new int[n];
    int mod = 0xffff;

    for (int i = 0; i < 2; i++) {
        int bucket[65536] = {0};
        int move = i * 16;
        for (int j = 0; j < n; j++) {
            bucket[(arr[j] >> move) & mod]++;
        }

        for (int j = 1; j < 65536; j++) {
            bucket[j] += bucket[j - 1];
        }

        for (int j = n - 1; j >= 0; j--) {
            backup[--bucket[(arr[j] >> move) & mod]] = arr[j];
        }

        swap(backup, arr);
    }

    delete[] backup;
}

在我的AMD Ryzen 7 5800X CPU上，当元素个数大于1000时，2048的基数排序速度就超过了系统排序，而且随着数组长度越来越大，系统排序和基数排序的速度差也越来越大。当数组长度小于10万时，我们可以选择基数2048，大于10万时，可以选择基数65536。具体怎么选择可以在自己的电脑上进行测试获得。

2.并行实现

基数排序的核心操作就是将n个数分配到不同的桶中，第一感觉是可以比较容易地实现并行化加速。但是仔细一想就会发现这里面的冲突非常严重。当不同的线程处理不同的数时，它们得到的基数可能相同，从而会同时去访问同一个桶，如果不加锁就会出错，加锁速度就会大幅降低。后来我搜到了一个openmp实现，然后适配上面的非并行2048基数排序，得到如下代码：

void parallel_radix_sort(vector<int> &arr) {
    int n = arr.size();
    vector<int> backup(n);
    int mod = 0x7ff;

    omp_set_num_threads(8); //需要根据自己的硬件进行调试

    for (int i = 0; i < 3; i++) {
        vector<int> bucket(2048);
        int local_bucket[2048] = {0};
        int move = i * 11;
#pragma omp parallel firstprivate(local_bucket)
        {
#pragma omp for schedule(static) nowait
            for (int j = 0; j < n; j++) {
                local_bucket[(arr[j] >> move) & mod]++;
            }

#pragma omp critical
            for (int j = 0; j < 2048; j++) {
                bucket[j] += local_bucket[j];
            }

#pragma omp barrier
#pragma omp single
            for (int j = 1; j < 2048; j++) {
                bucket[j] += bucket[j - 1];
            }

            int nthreads = omp_get_num_threads();
            int tid = omp_get_thread_num();
            for (int t = nthreads - 1; t >= 0; t--) {
                if (t == tid) {
                    for (int j = 0; j < 2048; j++) {
                        bucket[j] -= local_bucket[j];
                        local_bucket[j] = bucket[j];
                    }
                } else {
#pragma omp barrier
                }
            }

#pragma omp for schedule(static)
            for (int j = 0; j < n; j++) {
                backup[local_bucket[(arr[j] >> move) & mod]++] = arr[j];
            }
        }

        swap(backup, arr);
    }
}

上述代码和一般的openmp加速代码相比要复杂很多，在此给出一个简单的描述：为了规避不同线程在访问桶时存在的冲突，给每一个线程分配一个独享的桶，在上面代码中就是local_bucket。每个线程在算完自己的独享桶之后，再单线程汇总到共享桶中。最后每个线程再从共享桶中获取自己要写入的范围到local_bucket中进行写入。我曾经尝试对上述代码进行优化修改，但是任何的操作都会导致排序不正确，所以大家只需要拷贝使用即可。上述代码和非并行代码相比还有一个不同，backup用了vector而不是new，从我做实验的结果看差距不太大，new会略快，大家根据自己的编程习惯来选择即可。
上述并行代码用了8个线程来加速，这个在不同的CPU上可能不一样，需要大家做实验来验证多少线程最快。

3.基数排序的通用化

我们用基数排序实现了所有int型正整数的排序，如果数组中包含负整数时基数排序是否还可以使用？答案是肯定的。有两种解决方案：

1. 扫描整个数组获得最小值，如果最小值小于0，将每个元素减去最小值，然后做一次基数排序，完成之后再统一加上最小值。这种方案有个缺陷，当数据范围很大时，可能会存在溢出。但是在大多数情况下这种方案都是可行的；
2. 我们首先将整个数组按照正负号分成前后两部分，正整数部分正常做基数排序即可。负整数部分我们可以统一取绝对值，然后再做基数排序。做完基数排序之后再添加负号+reverse即可完成负整数的排序。这种方案也有缺陷，当数组中存在INT_MIN(-2^31)的时候也会存在溢出，但是概率要比方案一更小。可以通过一遍扫描将INT_MIN放在数组开头的位置来规避溢出的问题。我们在此给出方案二的代码实现：

//负数在前，正数在后
int move_negative(vector<int>& nums) {
    int n=nums.size();
    int i=0,j=n-1;
    while(i<=j) {
        while(i<n&&nums[i]<0) i++;
        while(j>=0&&nums[j]>=0) j--;
        if(i>=j) break;
        swap(nums[i++],nums[j--]);
    }

    return i;
}

void radix_sort(int *arr, int n) {
    int *backup = new int[n];
    int mod = 0x7ff;

    for (int i = 0; i < 3; i++) {
        int bucket[2048] = {0};
        int move = i * 11;
        for (int j = 0; j < n; j++) {
            bucket[(arr[j] >> move) & mod]++;
        }

        for (int j = 1; j < 2048; j++) {
            bucket[j] += bucket[j - 1];
        }

        for (int j = n - 1; j >= 0; j--) {
            backup[--bucket[(arr[j] >> move) & mod]] = arr[j];
        }

        swap(backup, arr);
    }

    swap(backup, arr);//需要再次交换，避免删除原始数组

    memcpy(arr, backup, n*sizeof(int)); //3次循环，需要再多一次拷贝

    delete[] backup;
}

void radix_sort(vector<int>& nums) {
	int pos=move_negative(nums);
	//负数取绝对值
    for(int i=0;i<pos;i++) {
        nums[i]=-nums[i];
    }

    radix_sort(nums.data(),pos);
    radix_sort(nums.data()+pos,nums.size()-pos);

	//反向+负号
    reverse(nums.begin(),nums.begin()+pos);
    for(int i=0;i<pos;i++) {
        nums[i]=-nums[i];
    }
}

除了包含正负整数的排序之外，我们还经常遇到float类型数组的排序。这种情况下基数排序依然可以使用。这是因为，如果我们忽略浮点数的格式，所有的浮点数按照大小排序之后，当做整数来看待也是排序的，所以我们可以直接将浮点数组转换成int型数组进行基数排序即可，具体实现就留给大家了。
最后，我们给出不同基数下的基数排序与系统排序的时间对比。可以看出，在10亿规模下，基数2048，线程数8下可以达到最快速度，大概耗时3.4s。和之前的N皇后问题一样，借用同事的13代i9之后耗时就可以降低到3s以内……

方法	1千万	1亿	10亿
系统排序	476ms	5.4s	60.5s
基数2048	92ms	0.94s	9.8s
基数65536	76ms	0.86s	9.5s
基数2048 8线程并行	31ms	0.36s	3.4s
基数65536 6线程并行	28ms	0.29s	3.9s