Python：更快的索引操作

Question

我有以下 sn-p，它使用规范索引提取序列中所有唯一值（可散列）的索引 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

Answer 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)}

Answer 3

我发现这个问题非常有趣，虽然我无法比其他提议的方法有很大的改进，但我确实找到了一种比其他提议的方法稍微快一点的纯 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 个不同的值时，我得到了：

1. get_values_and_indicies：每个循环 4.34 毫秒
2. data_to_idxlists_unique：每个循环 4.42 毫秒
3. by_pand2：每个循环 4.83 毫秒
4. data_to_idxlists：每个循环 6.09 毫秒
5. by_pand1：每个循环 9.39 毫秒
6. by_dd：每个循环 22.4 毫秒

如果我使用 10,000 个不同的值，我会得到：

1. get_values_and_indicies：每个循环 7.00 毫秒
2. data_to_idxlists_unique：每个循环 14.8 毫秒
3. by_dd：每个循环 29.8 毫秒
4. by_pand2：每个循环 47.7 毫秒
5. by_pand1：每个循环 67.3 毫秒
6. data_to_idxlists：每个循环 869 毫秒