Python:更快的索引操作

Question

我有以下片段,它提取序列中所有唯一值(hashable)的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

我发现这个问题非常有趣，虽然我无法比其他提出的方法获得很大的改进，但我确实找到了一个纯 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 毫秒