对列表中相同元素的索引进行分组的有效方法

Question

假设我有一个如下所示的列表：

[1, 2, 2, 5, 8, 3, 3, 9, 0, 1]

现在我想对相同元素的索引进行分组，所以结果应该如下所示：

[[0, 9], [1, 2], [3], [4], [5, 6], [7], [8]]

我如何以有效的方式做到这一点？我尽量避免使用循环，因此任何使用 numpy/pandas 函数的实现都很棒。

Answer 1

使用 pandas GroupBy.apply，这非常简单——使用您的数据对一系列索引进行分组。一个不错的好处是您可以保持索引的顺序。

data = [1, 2, 2, 5, 8, 3, 3, 9, 0, 1]
pd.Series(range(len(data))).groupby(data, sort=False).apply(list).tolist()
# [[0, 9], [1, 2], [3], [4], [5, 6], [7], [8]]

Answer 2

您可以使用 acollections.defaultdict对索引进行分组：

from collections import defaultdict

lst = [1, 2, 2, 5, 8, 3, 3, 9, 0, 1]

d = defaultdict(list)
for i, x in enumerate(lst):
    d[x].append(i)

print(list(d.values()))
# [[0, 9], [1, 2], [3], [4], [5, 6], [7], [8]]

它还保持添加索引的顺序而不进行排序。