[FIXED] Numpy array to dictionary with indices as values

Issue

I have a numpy array with integer values, let's call this array x.

I want to create some sort of list where for each value, I have the indices of x that hold this value.

For example, for:

x = [1,2,2,4,7,1,1,7,16]

I want to get:

{1: [0,5,6], 2:[1,2], 4:[3], 7:[4,7], 15:[16]}

The parenthesis I used are arbitrary, I don't care which data structure I use as long as I can output my result to a file as quickly as possible. At the end I want a .txt file that reads:

0,5,6

1,2

3

4,7

16

Solution

Since you mentioned you're not picky about the data structure of your values,tTo get something like the dictionary you posted in your question, you could do a dictionary comprehension over the unique values in x with np.where for the values:

>>> {i:np.where(x == i)[0] for i in set(x)}

{1: array([0, 5, 6]),
 2: array([1, 2]),
 4: array([3]),
 7: array([4, 7]),
 16: array([8])}

Comparing this to a more vanilla loop through a list, this will be significantly faster for larger arrays:

def list_method(x):
    res = {i:[] for i in set(x)}
    for i, value in enumerate(x):
         res[value].append(i)
    return res

def np_method(x):
    return {i:np.where(x == i)[0] for i in set(x)}

x = np.random.randint(1, 50, 1000000)


In [5]: %timeit list_method(x)
259 ms ± 4.03 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [6]: %timeit np_method(x)
120 ms ± 4.15 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Answered By - sacuL