user16054 user16054 - 4 months ago 32
Python Question

Python: How to index the elements of a numpy array?

I'm looking for a function that would do what the function

does in the following hypothetical code:

indices( numpy.array([[1, 2, 3], [2, 3, 4]]) )

{1: [(0,0)], 2: [(0,1),(1,0)], 3: [(0,2),(1,1)], 4: [(1,2)]}

Specifically, I want to produce a dictionary whose keys are the unique elements in the flattened array and whose values are lists of the full indices of the respective key.

I've looked at the
function, but it does not seem to provide an efficient way to solve this for large arrays. What's the best way to do this?

Notes: I'm using Python 2.7


Given that your desired output is a dictionary, I don't think there's going to be an efficient way to do this with NumPy operations. Your best bet will probably be something like

import collections
import itertools

d = collections.defaultdict(list)
for indices in itertools.product(*map(range, a.shape)):