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.
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 ;)
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]]
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)
