Efficient way to group indices of the same elements in a list

Question

Let's say I have a list that looks like:

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

Now I want to group the indices of the same elements, so the result should look like:

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

How do I do this in an efficient way? I try to avoid using loops so any implementations using numpy/pandas functions are great.

Answer 1

Using pandas GroupBy.apply, this is pretty straightforward—use your data to group on a Series of indices. A nice bonus here is you get to keep the order of your indices.

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

You can use a collections.defaultdict to group indices:

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

Which also maintains order of indices added without sorting.