Iterate over numpy with index (numpy equivalent of python enumerate)

Question

I\'m trying to create a function that will calculate the lattice distance (number of horizontal and vertical steps) between elements in a multi-dimensional numpy array. For

Answer 1

You can do it using np.ndenumerate but generally you don't need to iterate over an array.

You can simply create a meshgrid (or open grid) to get all indices at once and you can then process them (vectorized) much faster.

For example

>>> x, y = np.mgrid[slice(A.shape[0]), slice(A.shape[1])]
>>> x
array([[0, 0, 0],
       [1, 1, 1],
       [2, 2, 2]])
>>> y
array([[0, 1, 2],
       [0, 1, 2],
       [0, 1, 2]])

and these can be processed like any other array. So if your function that needs the indices can be vectorized you shouldn't do the manual loop!

For example to calculate the lattice distance for each point to a point say (2, 3):

>>> abs(x - 2) + abs(y - 3)
array([[5, 4, 3],
       [4, 3, 2],
       [3, 2, 1]])

For distances an ogrid would be faster. Just replace np.mgrid with np.ogrid:

>>> x, y = np.ogrid[slice(A.shape[0]), slice(A.shape[1])]
>>> np.hypot(x - 2, y - 3)  # cartesian distance this time! :-)
array([[ 3.60555128,  2.82842712,  2.23606798],
       [ 3.16227766,  2.23606798,  1.41421356],
       [ 3.        ,  2.        ,  1.        ]])

Answer 2

You can iterate through the values in your array with numpy.ndenumerate to get the indices of the values in your array.

Using the documentation above:

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
for index, values in np.ndenumerate(A):
    print(index, values)  # operate here

Answer 3

Another possible solution:

import numpy as np

A=np.array([[1,2,3],[4,5,6],[7,8,9]])
for _, val in np.ndenumerate(A):
    ind = np.argwhere(A==val)
    print val, ind

In this case you will obtain the array of indexes if value appears in array not once.