python - Numpy 数组作为查找表

Question

相关但不同的恕我直言:

(1) numpy: most efficient frequency counts for unique values in an array

(2) Using Numpy arrays as lookup tables

设置:

import numpy as np
from scipy.stats import itemfreq

x = np.array([1,  1,  1,  2,  25000, 2,  2,  5,  1,  1])
fq = itemfreq(x)
fq.astype(int)
array([[    1,     5],
       [    2,     3],
       [    5,     1],
       [25000,     1]])

现在，我想将 fq 用作查找表，然后执行以下操作:

res = magic_lookup_function(fq, x)
res
    array([5, 5, 5, 3, 1, 3, 3, 1, 5, 5])

如 (1) 和 (2) 中所建议的，我可以将 fq 转换为 python 字典，然后从那里进行查找，然后返回到 np.array。但是有没有更干净/更快/纯粹的 numpy 方法来做到这一点？

更新:此外，如 (2) 中所建议的，我可以使用 bincount，但我担心如果我的索引很大，例如~250,000。

谢谢!

更新的解决方案

正如@Jaime 指出的(下文)，np.unique 最多在 O(n log n) 时间内对数组进行排序。所以我想知道，itemfreq 背后发生了什么？结果是 itemfreq 对数组进行排序，我假设它也是 O(n log n):

In [875]: itemfreq??

def itemfreq(a):
... ... ...
    scores = _support.unique(a)
    scores = np.sort(scores)

这是一个timeit的例子

In [895]: import timeit

In [962]: timeit.timeit('fq = itemfreq(x)', setup='import numpy; from scipy.stats import itemfreq; x = numpy.array([ 1,  1,  1,  2, 250000,  2,  2,  5,  1,  1])', number=1000)
Out[962]: 0.3219749927520752

但是好像没必要对数组进行排序。如果我们在纯 Python 中执行此操作，则会发生以下情况。

In [963]: def test(arr):
   .....:     fd = {}
   .....:     for i in arr:
   .....:         fd[i] = fd.get(i,0) + 1
   .....:     return numpy.array([fd[j] for j in arr])

In [967]: timeit.timeit('test(x)', setup='import numpy; from __main__ import test; x = numpy.array([ 1,  1,  1,  2, 250000,  2,  2,  5,  1,  1])', number=1000)
Out[967]: 0.028257131576538086

哇，快 10 倍!

(至少，在这种情况下，数组不是太长，但可能包含大值。)

而且，仅供引用，正如我所怀疑的那样，使用 np.bincount 执行此操作对于大值来说效率低下:

In [970]: def test2(arr):
    bc = np.bincount(arr)
    return bc[arr]

In [971]: timeit.timeit('test2(x)', setup='import numpy; from __main__ import test2; x = numpy.array([ 1,  1,  1,  2, 250000,  2,  2,  5,  1,  1])', number=1000)
Out[971]: 0.0975029468536377

Answer 1

您可以使用 numpy.searchsorted :

def get_index(arr, val):                                                                
    index = np.searchsorted(arr, val)                                                            
    if arr[index] == val:                                                                        
        return index                                                                             

In [20]: arr = fq[:,:1].ravel()                                                                  

In [21]: arr
Out[21]: array([  1.,   2.,   5.,  25.])

In [22]: get_index(arr, 25)                                                                      
Out[22]: 3

In [23]: get_index(arr, 2)                                                                       
Out[23]: 1

In [24]: get_index(arr, 4)    #returns `None` for  item not found.