查找大鍵：字典與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

這裏的與的方法-

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

這假定lookup_array在其第一列排序鍵。如果不是這種情況，可以使用np.searchsorted的可選分揀機參數。

Answer 2

加載字典這巨大的內存有點不好，然後增加查找的開銷。如果這是一種很常用的數據結構，那麼如何使用數據庫引擎。如果你不喜歡SQL，那麼有KEY/VALUE數據庫。它們針對查找進行了高度優化。

Answer 3

如果某些特殊情況適用，您可以使用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