Create dictionary of value frequencies from numpy array

Question

I need to create a dictionary from a loop that runs through an array of 2 columns of numbers. Below is a small subset of the array:

array([[  0,   1],
       [  1,   0],
       [  1,   2],
       [  2,   3],
       [  2,   1]])

I would like to create a dictionary that takes the unique numbers of the first column as keys (e.g. {0,1,2} in this example) and the corresponding numbers in the second column as values.

For this example the dictionary would look like:

dict = {0:[1], 1:[0,2], 2:[3,1]}

My array is very long (370,000 x 2) so I would like to do this through an efficient loop. Any advice would be greatly appreciated!

Answer 1

You can use defaultdict to accomplish this.

from collections import defaultdict
a = np.array([[  0,   1],[  1,   0],[  1,   2],[  2,   3], [  2,   1]])
d = defaultdict(list)
for x,y in a:
    d[x].append(y)