How to extract the indices of equal elements from a list?

user_5420863 0
Viewed 108

If my input a = [3.0, 4.0, 2.5, 2.5, 3.0], There are two 3.0, one 4.0, and two 2.5

As you can see, a[0] = a[4], a[1] is just by itself, and a[2] = a[3]

I want my output as [[0,4], [1], [2,3]]

How would I code this? Is there a built-in function to do this? If not, any hint or suggestion? Thank you.

numpy-ndarray python-3.x list numpy python
3 Answers

First, a simple solution leveraging collections.defaultdict. Be careful, this only works for flat (1-dimensional) lists! 

from collections import defaultdict 

def group_by_value(lst_in):
    res_dict = defaultdict(list)
    for idx, item in enumerate(lst_in):
        res_dict[item].append(idx)
    return res_dict

defaultdict is just like a dict except that when you initialize it like d = defaultdict(list) then d[x] would resolve to an empty list [] if x is not among the keys defined in d.

Here is a solution using numpy which works regardless of the number of dimensions. Credit for the idea goes to @meTchaikovsky.

import numpy as np

def group_by_value(arr):
    return {elem: list(zip(*np.where(arr == elem))) for elem in np.unique(arr)}

Thanks to @AlexanderCécile for refactoring this answer ;)

user_51200890 
sol = {}        
for i, j in enumerate(a):
    if j in sol:
        sol[j].append(i)
    else:
        sol[j] = [i]
[sol[i] for i in sol]
[[0, 4], [1], [2, 3]]
user_73694310

You can do this with the help of numpy

import numpy as np

a = np.array(a)
result = [np.arange(len(a))[a == item].tolist() for item in np.unique(a)]

in which result is the list you want.

To do it without using numpy

result = []
unique_vals = set(a)

for val in unique_vals:
    tmp = []
    for ind, item in enumerate(a):
        if item == val:
            tmp.append(ind)
    result.append(tmp)
user_83668050
Related