我有以下代码段，它使用规范索引在类似序列的data中提取所有唯一值（可散列）的索引，并将它们作为列表存储在字典中：
from collections import defaultdict
idx_lists = defaultdict(list)
for idx, ele in enumerate(data):
idx_lists[ele].append(idx)
这对我来说是一个非常常见的用例。而且我的代码执行时间的90％花在了这几行上。这部分在执行期间传递了10000多次，每次运行时len(data)约为50000到100000。唯一元素的数量大致为50到150。
是否有更快的方法，可能是矢量化/ c扩展（例如numpy或pandas方法），它实现了相同的目标？
非常感谢。
答案 0 :(得分：5)
没有我原先希望的那么令人印象深刻（在groupby代码路径中仍有相当多的纯Python），但你可能能够将时间减少2-4倍，具体取决于你的数量关心所涉及的确切最终类型：
>>> %timeit by_dd(data)
10 loops, best of 3: 42.9 ms per loop
>>> %timeit by_pand1(data)
100 loops, best of 3: 18.2 ms per loop
>>> %timeit by_pand2(data)
100 loops, best of 3: 11.5 ms per loop
给了我
if __name__ == '__main__':
country = input('\nPlease enter the country> ')
答案 1 :(得分：2)
虽然它不是完美的解决方案（它是O（NlogN）而不是O（N）），但更快，矢量化的方法是：
def data_to_idxlists(data):
sorting_ixs = np.argsort(data)
uniques, unique_indices = np.unique(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])}
另一种解决方案是O（N * U），（其中U是唯一组的数量）：
def data_to_idxlists(data):
u, ixs = np.unique(data, return_inverse=True)
return {u: np.nonzero(ixs==i) for i, u in enumerate(u)}
答案 2 :(得分：1)
我发现这个问题非常有趣，虽然我无法对其他提议的方法进行大的改进，但我确实找到了一种比其他方法略快的纯粹numpy方法。
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，则可以将最后一次更改为return defaultdict(list, ((unique_vals[i], sorted_indices[run_endpoints[i]:run_endpoints[i + 1]].tolist()) for i in range(num_values)))，这会将我的测试中的总体时间增加到4.4毫秒。
最后，我应该注意这些时间对数据敏感。当我只使用10个不同的值时：
如果我使用了10,000个不同的值：