from collections import defaultdict
idx_lists = defaultdict(list)
for idx,ele in enumerate(data):
idx_lists[ele].append(idx)
这对我来说是一个很常见的用例.而且我的代码执行时间的90％花在了这几行上.该部分在执行期间传递超过10000次,并且每次运行时len(数据)大约为50000到100000.独特元素的数量大致为50到150.
有没有更快的方法,也许是矢量化/ c扩展(例如numpy或pandas方法),实现了同样的目的？
非常感谢.
解决方法
import numpy as np
import pandas as pd
from collections import defaultdict
data = np.random.randint(0,10**2,size=10**5)
series = pd.Series(data)
def get_values_and_indicies(input_data):
input_data = np.asarray(input_data)
sorted_indices = input_data.argsort() # Get the sorted indices
# Get the sorted data so we can see where the values change
sorted_data = input_data[sorted_indices]
# Find the locations where the values change and include the first and last values
run_endpoints = np.concatenate(([0],np.where(sorted_data[1:] != sorted_data[:-1])[0] + 1,[len(input_data)]))
# Get the unique values themselves
unique_vals = sorted_data[run_endpoints[:-1]]
# Return the unique values along with the indices associated with that value
return {unique_vals[i]: sorted_indices[run_endpoints[i]:run_endpoints[i + 1]].tolist() for i in range(num_values)}
def by_dd(input_data):
idx_lists = defaultdict(list)
for idx,ele in enumerate(input_data):
idx_lists[ele].append(idx)
return idx_lists
def by_pand1(input_data):
idx_lists = defaultdict(list)
return {k: v.tolist() for k,v in series.groupby(input_data).indices.items()}
def by_pand2(input_data):
return series.groupby(input_data).indices
def data_to_idxlists(input_data):
u,ixs = np.unique(input_data,return_inverse=True)
return {u: np.nonzero(ixs==i) for i,u in enumerate(u)}
def data_to_idxlists_unique(input_data):
sorting_ixs = np.argsort(input_data)
uniques,unique_indices = np.unique(input_data[sorting_ixs],return_index = True)
return {u: sorting_ixs[start:stop] for u,start,stop in zip(uniques,unique_indices,list(unique_indices[1:])+[None])}
由此产生的时间(从最快到最慢)：
>>> %timeit get_values_and_indicies(data) 100 loops,best of 3: 4.25 ms per loop >>> %timeit by_pand2(series) 100 loops,best of 3: 5.22 ms per loop >>> %timeit data_to_idxlists_unique(data) 100 loops,best of 3: 6.23 ms per loop >>> %timeit by_pand1(series) 100 loops,best of 3: 10.2 ms per loop >>> %timeit data_to_idxlists(data) 100 loops,best of 3: 15.5 ms per loop >>> %timeit by_dd(data) 10 loops,best of 3: 21.4 ms per loop
应该注意的是,与by_pand2不同,它会产生一个列表的字典,如示例中给出的.如果你更愿意返回一个defaultdict,你可以简单地改变最后一次返回defaultdict(list,((unique_vals [i],sorted_indices [run_endpoints [i]：run_endpoints [i 1]].tolist())for i in range (num_values)))这使我的测试中的总体时间增加到4.4毫秒.
最后,我应该注意到这些时间对数据敏感.当我使用10个不同的值时,我得到了：
> get_values_and_indicies：每个循环4.34毫秒
> data_to_idxlists_unique：每个循环4.42 ms
> by_pand2：每循环4.83 ms
> data_to_idxlists：每个循环6.09毫秒
> by_pand1：每个循环9.39 ms
> by_dd：每个循环22.4毫秒
如果我使用了10,000个不同的值,我得到了：
> get_values_and_indicies：每个循环7.00毫秒> data_to_idxlists_unique：每个循环14.8 ms> by_dd：每个循环29.8 ms> by_pand2：每个循环47.7毫秒> by_pand1：每个循环67.3毫秒> data_to_idxlists：每个循环869 ms