python - 查找大量键 : dictionary vs. 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 索引作为字典查找的非常快速的替代方法。

• 键必须是整数

• 您有足够的内存来创建大小与您希望查找的最大键值(以便所有键都对应于数组中的有效索引。)

思路是用

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