Numpy 数组作为查找表答案

【问题标题】：Numpy array as lookup tableNumpy 数组作为查找表
【发布时间】：2014-03-25 06:59:15
【问题描述】：

设置：

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

【问题讨论】：

标签： python numpy

【解决方案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.

【讨论】：

啊，很好。虽然，searchsorted 的效率如何？在我看来，充其量是 O(log n) （但我在这里可能是错的）。理想情况下，我希望找到像 python 字典查找这样 O(1) 的东西，但没有来回转换的开销。（这可能是不可能的......）
@andros 是的，这是O(log N)，我忘了说它还期望数据被排序。
@andros 但是请注意，在O(log N) 中，搜索是由 C 类型上的高效 C 代码完成的。即使N 的值非常大，如果searchsorted 花费的时间与字典查找的时间大致相同，这并不奇怪。我会考虑尝试并将其作为一种可能的解决方案进行分析，即使它渐进地更糟，因为对于您甚至不考虑的 N 这么大的值，它可能会变得更糟。

【解决方案2】：

因为您的查找表不仅仅是任何查找表，而是频率列表，您可能需要考虑以下选项：

>>> x = np.array([1,  1,  1,  2,  25, 2,  2,  5,  1,  1])
>>> x_unq, x_idx = np.unique(x, return_inverse=True)
>>> np.take(np.bincount(x_idx), x_idx)
array([5, 5, 5, 3, 1, 3, 3, 1, 5, 5], dtype=int64)

即使您的查找表更复杂，即：

>>> lut = np.array([[ 1, 10],
...                 [ 2,  9],
...                 [ 5,  8],
...                 [25,  7]])

如果您负担得起使用return_index 执行np.unique（它对数组进行排序，因此它是n log n）调用，您可以使用小的连续整数作为索引，这通常会使事情变得更容易，例如使用np.searchsorted 你会这样做：

>>> np.take(lut[:, 1], np.take(np.searchsorted(lut[:, 0], x_unq), x_idx))
array([10, 10, 10,  9,  7,  9,  9,  8, 10, 10])

【讨论】：