python - 如何比快速排序更快地对整数数组进行排序？

Question

使用 numpy 的快速排序对整数数组进行排序已成为 我算法的瓶颈。不幸的是，numpy 没有 radix sort yet . 虽然counting sort将是 numpy 中的一行:

np.repeat(np.arange(1+x.max()), np.bincount(x))

请参阅 How can I vectorize this python count sort so it is absolutely as fast as it can be? 的已接受答案问题，整数 在我的应用程序中可以从 0 运行到 2**32。

我被快速排序困住了吗？

---

<子> 这篇文章的主要动机是 [Numpy 分组使用 itertools.groupby 性能](https://stackoverflow.com/q/4651683/341970) 问题。 还要注意的是 [不仅可以提出和回答您自己的问题，还明确鼓励这样做。](https://blog.stackoverflow.com/2011/07/its-ok-to-ask-and-answer-your-own -问题/)

Answer 1

不，您没有被快速排序困住。例如，您可以使用 integer_sort 来自 Boost.Sort 或来自 usort 的 u4_sort .对该数组进行排序时:

array(randint(0, high=1<<32, size=10**8), uint32)

我得到以下结果:

NumPy quicksort:         8.636 s  1.0  (baseline)
Boost.Sort integer_sort: 4.327 s  2.0x speedup
usort u4_sort:           2.065 s  4.2x speedup

I would not jump to conclusions based on this single experiment and use usort blindly. I would test with my actual data and measure what happens. Your mileage will vary depending on your data and on your machine. The integer_sort in Boost.Sort has a rich set of options for tuning, see the documentation.

Below I describe two ways to call a native C or C++ function from Python. Despite the long description, it's fairly easy to do it.

---

Boost.Sort

Put these lines into the spreadsort.cpp file:

#include <cinttypes>
#include "boost/sort/spreadsort/spreadsort.hpp"
using namespace boost::sort::spreadsort;

extern "C" {
    void spreadsort(std::uint32_t* begin,  std::size_t len) {
        integer_sort(begin, begin + len);
    }
}

它基本上实例化了 32 位的模板化 integer_sort 无符号整数； extern "C" 部分通过禁用来确保 C 链接 名称修改。 假设您正在使用 gcc 并且必要的包含 boost 文件 在/tmp/boost_1_60_0目录下，编译即可:

g++ -O3 -std=c++11 -march=native -DNDEBUG -shared -fPIC -I/tmp/boost_1_60_0 spreadsort.cpp -o spreadsort.so  

关键标志是-fPIC来生成 position-independet code 和 -shared 生成一个 shared object .so 文件。 (阅读 gcc 的文档以获取更多详细信息。)

然后，包装 spreadsort() C++ 函数 在 Python 中使用 ctypes :

from ctypes import cdll, c_size_t, c_uint32
from numpy import uint32
from numpy.ctypeslib import ndpointer

__all__ = ['integer_sort']

# In spreadsort.cpp: void spreadsort(std::uint32_t* begin,  std::size_t len)
lib = cdll.LoadLibrary('./spreadsort.so')
sort = lib.spreadsort
sort.restype = None
sort.argtypes = [ndpointer(c_uint32, flags='C_CONTIGUOUS'), c_size_t]

def integer_sort(arr):
    assert arr.dtype == uint32, 'Expected uint32, got {}'.format(arr.dtype)
    sort(arr, arr.size)

或者，您可以使用 cffi :

from cffi import FFI
from numpy import uint32

__all__ = ['integer_sort']

ffi = FFI()
ffi.cdef('void spreadsort(uint32_t* begin,  size_t len);')
C = ffi.dlopen('./spreadsort.so')

def integer_sort(arr):
    assert arr.dtype == uint32, 'Expected uint32, got {}'.format(arr.dtype)
    begin = ffi.cast('uint32_t*', arr.ctypes.data)
    C.spreadsort(begin, arr.size)

在 cdll.LoadLibrary() 和 ffi.dlopen() 调用中，我假设 spreadsort.so 文件的路径是 ./spreadsort.so。或者， 你可以写

lib = cdll.LoadLibrary('spreadsort.so')

或

C = ffi.dlopen('spreadsort.so')

如果将 spreadsort.so 的路径附加到 LD_LIBRARY_PATH 环境 多变的。另见 Shared Libraries .

用法。在这两种情况下，您只需调用上面的 Python 包装函数 integer_sort() 使用 32 位无符号整数的 numpy 数组。

---

排序

对于u4_sort，可以这样编译:

cc -DBUILDING_u4_sort -I/usr/include -I./ -I../ -I../../ -I../../../ -I../../../../ -std=c99 -fgnu89-inline -O3 -g -fPIC -shared -march=native u4_sort.c -o u4_sort.so

在 u4_sort.c 文件所在的目录中发出此命令。 (可能有一种不那么骇人听闻的方式，但我没能弄明白。我 只是查看了 usort 目录中的 deps.mk 文件以找出 必要的编译器标志和包含路径。)

然后，您可以按如下方式包装 C 函数:

from cffi import FFI
from numpy import uint32

__all__ = ['integer_sort']

ffi = FFI()
ffi.cdef('void u4_sort(unsigned* a, const long sz);')
C = ffi.dlopen('u4_sort.so')

def integer_sort(arr):
    assert arr.dtype == uint32, 'Expected uint32, got {}'.format(arr.dtype)
    begin = ffi.cast('unsigned*', arr.ctypes.data)
    C.u4_sort(begin, arr.size)

在上面的代码中，我假设到 u4_sort.so 的路径是 附加到 LD_LIBRARY_PATH 环境变量。

用法。与之前使用 Boost.Sort 一样，您只需使用 32 位无符号整数的 numpy 数组调用上述 Python 包装函数 integer_sort()。