numpy.searchsorted 在结构化数组上的性能很差

Question

_{如果我误用了任何术语，请提前道歉，请随时更正。}

我有一个带有dtype '<f16, |S30' 的排序数组。当我在第一个字段上使用searchsorted 时，它的运行速度非常慢（300 万个项目大约需要 0.4 秒）。这比bisect 在一个普通的 Python 元组列表上做同样的事情要长得多。

%timeit a['f0'].searchsorted(400.)
1 loops, best of 3: 398 ms per loop

但是，如果我将浮点部分复制到另一个单独的数组中，则搜索速度比 bisect 快：

b = a['f0'].copy()

%timeit b.searchsorted(400.)
1000000 loops, best of 3: 945 ns per loop

我的问题是：

1. 是我做错了什么还是 NumPy 中的回归？
2. 有没有办法在不重复数据的情况下避免这种情况？

Answer 1

我记得前段时间看到过这个。如果我没记错的话，我认为 searchsorted 在数据不连续时会制作数据的临时副本。如果稍后我有时间，我会查看代码以确认正在发生的事情（或者也许更熟悉代码的人可以确认这一点）。

同时，如果您不想重组代码以避免使用结构化数组，那么最好的选择可能是使用bisect_left(a['f0'], 400.)。在我的机器上，它比在连续数组上搜索排序慢 8 倍，但比在非连续数组上搜索排序快 1000 倍。

In [5]: a = np.arange((6e6)).view([('f0', float), ('f1', float)])

In [6]: timeit a['f0'].searchsorted(400.)
10 loops, best of 3: 51.1 ms per loop

In [7]: timeit a['f0'].copy()
10 loops, best of 3: 51 ms per loop

In [8]: timeit bisect_left(a['f0'], 400.)
10000 loops, best of 3: 52.8 us per loop

In [9]: f0 = a['f0'].copy()

In [10]: timeit f0.searchsorted(400.)
100000 loops, best of 3: 7.85 us per loop

Answer 2

这里有一些代码来说明问题的规模（截至 2015 年 5 月 11 日）以及如何“修复”它。

import numpy as np
import bisect
import timeit
from random import randint

dtype = np.dtype([ ('pos','<I'),('sig','<H') ])             # my data is unsigned 32bit, and unsigned 16bit
data1 = np.fromfile('./all2/840d.0a9b45e8c5344abf6ac761017e93b5bb.2.1bp.binary', dtype)

dtype2 = np.dtype([('pos',np.uint32),('sig',np.uint32)])    # convert data to both unsigned 32bit
data2 = data1.astype(dtype2)

data3 = data2.view(('uint32', len(data2.dtype.names)))    # convert above to a normal array (not structured array)

print data1.dtype.descr # [('pos', '<u4'), ('sig', '<u2')]
print data2.dtype.descr # [('pos', '<u4'), ('sig', '<u4')]
print data3.dtype.descr # [('', '<u4')]

print data1.nbytes  # 50344494
print data2.nbytes  # 67125992
print data3.nbytes  # 67125992

print data1['pos'].max() # 2099257376
print data2['pos'].max() # 2099257376
print data3[:,0].max()   # 2099257376

def b1():   return bisect.bisect_left(data1['pos'],           randint(100000000,200000000))
def b2():   return bisect.bisect_left(data2['pos'],           randint(100000000,200000000))
def b3():   return bisect.bisect_left(data3[:,0],             randint(100000000,200000000))
def ss1():  return np.searchsorted(data1['pos'],              randint(100000000,200000000))
def ss2():  return np.searchsorted(data2['pos'],              randint(100000000,200000000))
def ss3():  return np.searchsorted(data3[:,0],                randint(100000000,200000000))

def ricob1():   return bisect.bisect_left(data1['pos'], np.uint32(randint(100000000,200000000)))
def ricob2():   return bisect.bisect_left(data2['pos'], np.uint32(randint(100000000,200000000)))
def ricob3():   return bisect.bisect_left(data3[:,0],   np.uint32(randint(100000000,200000000)))
def ricoss1():  return np.searchsorted(data1['pos'],    np.uint32(randint(100000000,200000000)))
def ricoss2():  return np.searchsorted(data2['pos'],    np.uint32(randint(100000000,200000000)))
def ricoss3():  return np.searchsorted(data3[:,0],      np.uint32(randint(100000000,200000000)))

print timeit.timeit(b1,number=300)  # 0.0085117816925
print timeit.timeit(b2,number=300)  # 0.00826191902161
print timeit.timeit(b3,number=300)  # 0.00828003883362
print timeit.timeit(ss1,number=300) # 6.57477498055
print timeit.timeit(ss2,number=300) # 5.95308589935
print timeit.timeit(ss3,number=300) # 5.92483091354

print timeit.timeit(ricob1,number=300)  # 0.00120902061462
print timeit.timeit(ricob2,number=300)  # 0.00120401382446
print timeit.timeit(ricob3,number=300)  # 0.00120711326599
print timeit.timeit(ricoss1,number=300) # 4.39265394211
print timeit.timeit(ricoss2,number=300) # 0.00116586685181
print timeit.timeit(ricoss3,number=300) # 0.00108909606934

更新！ 因此，多亏了 Rico 的 cmets，似乎为您要搜索的数字设置类型排序/平分确实很重要！ 但是，在具有 32 位和 16 位整数的结构化数组上，它仍然很慢（尽管没有以前那么慢了）