Python:更快的索引操作

Question

我有以下代码段，它提取具有规范索引的类序列 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 方法)，实现同样的事情？

非常感谢。

Answer 1

不像我最初希望的那样令人印象深刻(在 groupby 代码路径中仍然有相当多的纯 Python)，但是你可以将时间减少 2-4 倍，这取决于你有多少关心涉及的确切最终类型:

import numpy as np, pandas as pd
from collections import defaultdict

def by_dd(data):
    idx_lists = defaultdict(list)
    for idx, ele in enumerate(data):
        idx_lists[ele].append(idx)
    return idx_lists

def by_pand1(data):
    return {k: v.tolist() for k,v in data.groupby(data.values).indices.items()}

def by_pand2(data):
    return data.groupby(data.values).indices

data = pd.Series(np.random.randint(0, 100, size=10**5))    

给我

>>> %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