Numpy array to dictionary with indices as values

Question

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

Answer 1

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)

Answer 2

Pure python will be like this:

result = {}
for idx,val in enumerate(x):
    arr = result.get(val,[])
    arr.append(idx)
    result[val] = arr

Answer 3

x = [1,2,2,4,7,1,1,7,16]
numlist = []
numdict = {}
c = 0
for n in x:
    if n not in numlist:
        numlist.append(n)
        numdict[n] = [c]
    else:
        numdict[n].append(c)
    c += 1
print(numlist, numdict)

Output is: [1, 2, 4, 7, 16] {1: [0, 5, 6], 2: [1, 2], 4: [3], 7: [4, 7], 16: [8]} To write to file use:

with open('file.txt', 'w') as f:
    f.write(str(numdict))