使用键列表从字典中快速查找字典

Question

我的字典：

d = {'a':1, 'b':2, 'c':3}

和我的钥匙清单：

keys = np.array(['a','b','a','c','a','b'])

我想在不使用 for 循环的情况下获得相应值的列表

我以以下方式尝试使用 for 循环，但对于我的工作目的而言，它的计算成本太高了。

这是for循环版本。

l = [d[i] for i in keys]

你知道没有 FOR LOOPS 的版本吗，也许利用 np.array 的广播..掩码？

Answer 1

对于一般情况，列表理解方法[d[i] for i in keys]很好。

---

对于非常大的列表，一种提高性能的方法是定义一个结构化数组，它允许使用混合类型，并使用np.searchsorted ：

def str_array(d, keys):
    items = list(d.items())
    k, v = zip(*items)
    dtype_v = np.max(v).itemsize
    dtype_k = np.array(k).dtype
    a = np.array(items, dtype=[('key', dtype_k), 
                               ('value', f'i{dtype_v}')])
    ixs_s = np.argsort(a['key'])
    k_ixs = ixs_s[np.searchsorted(a['key'], keys, sorter=ixs_s)]
    return a['value'][k_ixs]

str_array(d,keys)
# array([1, 2, 1, 3, 1, 2])

---

让我们将性能与其他一些典型方法进行比较：

d = {'key1':100, 'some_other_key':8, 'key3':15, 'nth_key':0}

perfplot.show(
    setup=lambda n: np.random.choice(list(d.keys()), size=n), 

    kernels=[
        lambda x: np.array([d[i] for i in x]),
        lambda x: np.vectorize(d.get)(x),
        lambda x: pd.Series(d).loc[x].values,
        lambda x: operator.itemgetter(*x)(d),
        lambda x: str_array(d, x),
    ],

    labels=['list-comp', 'np.vectorize', 'pd.loc', 'itemgetter', 'str_array'],
    n_range=[2**k for k in range(0, 20)],
    xlabel='N'
)

---

因此，例如对于n=100_000 ：

keys = np.random.choice(list(d.keys()), size=100_000)

%timeit str_array(d, keys)
# 5.51 ms ± 231 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit [d[i] for i in keys]
# 51.7 ms ± 1.97 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

使用np.searchsorted的第二种方法比使用简单的列表理解快 10 倍。

Answer 2

我不知道相对性能，但我发现这个解决方案非常快速和简单。 将您的键转换为一个系列，然后使用内置的 map 函数来返回您的答案。

import pandas as pd
d = {'a':1, 'b':2, 'c':3}
keys = np.array(['a','b','a','c','a','b'])
keys1 = pd.Series(keys)
keys1.map(d)