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

【问题标题】：Looking up large sets of keys: dictionary vs. NumPy array查找大量键：字典与 NumPy 数组
【发布时间】：2016-08-07 17:52:06
【问题描述】：

我有一组非常大（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]

但这似乎也很耗时。

有没有其他方法可以在不迭代的情况下快速查找大量不连续的值？

【问题讨论】：

在第一个示例中查找是字典
不应该是lookup_array[:,0]吗？另外，key_set 是否包含 unique 键？
是的，键是整数。
key_set 是否包含不在lookup_array 中的值？
不，key_set 是lookup_array[:,0] 的子集。不，每个键只有一个值。

标签： python arrays numpy dictionary iteration

【解决方案1】：

如果适用某些特殊条件，您可以使用 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

【讨论】：

【解决方案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 一起使用。

【讨论】：

【解决方案3】：

在内存中加载这么大的字典有点不好，然后会增加查找的开销。如果这是您经常使用的数据结构，那么如何使用数据库引擎。如果您不喜欢 SQL，则可以使用 KEY / VALUE 数据库。它们针对查找进行了高度优化。

【讨论】：