Question

使用numpy的quicksort对整数数组进行排序已经成为了我算法的瓶颈。不幸的是，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。

我坚持使用quicksort吗？

<子> 这篇文章的主要动机是 Numpy grouping using itertools.groupby performance 问题。
另请注意 it is not merely OK to ask and answer your own question, it is explicitly encouraged.

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

我不会根据这个单一的实验和使用来得出结论 {盲目地usort我会测试我的实际数据并测量会发生什么。您的里程将会因您的数据和计算机而异。该 Boost.Sort中的integer_sort有一组丰富的调优选项，请参阅 documentation

下面我将介绍两种从Python调用本机C或C ++函数的方法。尽管有很长的描述，但它很容易实现。

<强> Boost.Sort

将这些行放入spreadort.cpp文件中：

#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数组。

<强> usort

对于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()。

Answer 2

基数排序基数256（1字节）可以生成一个计数矩阵，用于在1次传递数据时确定存储桶大小，然后需要4次传递来进行排序。在我的系统上，Intel 2600K 3.4ghz，使用Visual Studio版本构建C ++程序，需要大约1.15秒来对10 ^ 8个伪随机无符号32位整数进行排序。

查看Ali回答中提到的u4_sort代码，程序类似，但是u4_sort使用{10,11,11}的字段大小，需要3次传递来排序数据，1次传递要复制回来，而此示例使用字段大小{8,8,8,8}，需要4次传递才能对数据进行排序。由于随机访问写入，该过程可能是内存带宽受限，因此u4_sort中的优化，如移位宏，每个字段固定移位的一个循环，并没有多大帮助。我的结果可能更好，可能是由于系统和/或编译器的差异。（注意u8_sort用于64位整数）。

示例代码：

//  a is input array, b is working array
void RadixSort(uint32_t * a, uint32_t *b, size_t count)
{
size_t mIndex[4][256] = {0};            // count / index matrix
size_t i,j,m,n;
uint32_t u;
    for(i = 0; i < count; i++){         // generate histograms
        u = a[i];
        for(j = 0; j < 4; j++){
            mIndex[j][(size_t)(u & 0xff)]++;
            u >>= 8;
        }       
    }
    for(j = 0; j < 4; j++){             // convert to indices
        m = 0;
        for(i = 0; i < 256; i++){
            n = mIndex[j][i];
            mIndex[j][i] = m;
            m += n;
        }       
    }
    for(j = 0; j < 4; j++){             // radix sort
        for(i = 0; i < count; i++){     //  sort by current lsb
            u = a[i];
            m = (size_t)(u>>(j<<3))&0xff;
            b[mIndex[j][m]++] = u;
        }
        std::swap(a, b);                //  swap ptrs
    }
}

Answer 3

根据@rcgldr C程序，使用python/numba(0.23)进行基数排序，在2核处理器上使用多线程。

首先在numba上进行基数排序，有两个全局数组以提高效率。

from threading import Thread
from pylab import *
from numba import jit
n=uint32(10**8)
m=n//2
if 'x1'  not in locals() : x1=array(randint(0,1<<16,2*n),uint16); #to avoid regeneration
x2=x1.copy()
x=frombuffer(x2,uint32) # randint doesn't work with 32 bits here :(
y=empty_like(x) 
nbits=8
buffsize=1<<nbits
mask=buffsize-1

@jit(nopython=True,nogil=True)
def radix(x,y):
    xs=x.size
    dec=0
    while dec < 32 :
        u=np.zeros(buffsize,uint32)
        k=0
        while k<xs:
            u[(x[k]>>dec)& mask]+=1
            k+=1
        j=t=0
        for j in range(buffsize):
            b=u[j]
            u[j]=t
            t+=b
            v=u.copy()
        k=0
        while k<xs:
            j=(x[k]>>dec)&mask
            y[u[j]]=x[k]
            u[j]+=1
            k+=1
        x,y=y,x
        dec+=nbits

然后是parallélisation，可以使用nogil选项。

def para(nthreads=2):
        threads=[Thread(target=radix,
            args=(x[i*n//nthreads(i+1)*n//nthreads],
            y[i*n//nthreads:(i+1)*n//nthreads])) 
            for i in range(nthreads)]
        for t in  threads: t.start()
        for t in  threads: t.join()

@jit
def fuse(x,y):
    kl=0
    kr=n//2
    k=0
    while k<n:
        if y[kl]<x[kr] :
            x[k]=y[kl]
            kl+=1
            if kl==m : break
        else :
            x[k]=x[kr]
            kr+=1
        k+=1

def sort():
    para(2)
    y[:m]=x[:m]
    fuse(x,y)

基准：

In [24]: %timeit x2=x1.copy();x=frombuffer(x2,uint32) # time offset
1 loops, best of 3: 431 ms per loop

In [25]: %timeit x2=x1.copy();x=frombuffer(x2,uint32);x.sort()
1 loops, best of 3: 37.8 s per loop

In [26]: %timeit x2=x1.copy();x=frombuffer(x2,uint32);para(1)
1 loops, best of 3: 5.7 s per loop

In [27]: %timeit x2=x1.copy();x=frombuffer(x2,uint32);sort()
1 loops, best of 3: 4.02 s per loop

这是一款纯粹的python解决方案，在我可怜的1GHz机器上获得10倍（37s-> 3.5s）的增益。可以通过更多核心和multifusion进行增强。

如何比quicksort更快地对整数数组进行排序？

3 个答案: