查找大量键：字典与NumPy数组

Question

我有很大的一组键/值对（200k +），为此我需要检索非常大（有时是全部）的值。显而易见的方法是使用这样的字典：

 values = {lookup.get(key) for key in key_set}

这在我的代码中变得非常耗时，并且我想知道是否存在一种更快的方法来使用NumPy数组来实现。我一直在尝试使用具有两列和n行的数组，这样对于任何单个键：

value = lookup_array[lookup_array[:,0] == key, 1]

但是我不确定如何在不进行昂贵的迭代的情况下将其扩展到很多键。我看了看：

values = lookup_array[np.in1d(lookup_array[:,0], key_set), 1]

但这似乎也很耗时。

还有其他方法可以快速进行大量非连续值的查找而不进行迭代吗？

Answer 1

如果某些特殊条件适用，您可以使用 NumPy 索引作为字典查找的非常快速的替代方法。

• 键必须是整数

• 您有足够的内存来创建一个 NumPy 数组，其大小与您希望查找的最大键值一样大（以便所有键都对应于数组中的有效索引。）

这个想法是使用

lookup_array = np.empty((M,), dtype=values.dtype)
lookup_array[keys] = values
result = lookup_array[key_set]

代替

result = {lookup_dict.get(key) for key in key_set}

例如，

import numpy as np
import pandas as pd

def using_dict(lookup_dict, key_set):
    return {lookup_dict.get(key) for key in key_set}

def using_array(lookup_array, key_set):
    return lookup_array[key_set]

def using_pandas(df, key_set):
    return df.loc[df['a'].isin(key_set)]

M = 10**6
N = 2*10**5
K = 10**4
keys = np.random.randint(M, size=(N,))
values = np.random.random((N,))
lookup_dict = dict(zip(keys, values))
lookup_array = np.empty((M,), dtype=values.dtype)
lookup_array[keys] = values
df = pd.DataFrame(np.column_stack([keys, values]), columns=list('ab'))
key_set = np.random.choice(keys, size=(K,))

这是上述方法的 timeit 基准测试（使用 IPython）：

In [25]: %timeit using_array(lookup_array, key_set)
10000 loops, best of 3: 22.4 µs per loop

In [26]: %timeit using_dict(lookup_dict, key_set)
100 loops, best of 3: 3.73 ms per loop

In [24]: %timeit using_pandas(df, key_set)
10 loops, best of 3: 38.9 ms per loop

Answer 2

这是一种方法np.searchsorted-

row_idx = np.searchsorted(lookup_array[:,0],key_set)[key_set.argsort()]
values = lookup_array[row_idx,1]

这假设lookup_array将键排序在其第一列中。如果不是这种情况，您可以将可选的 sorter 参数与np.searchsorted.