Indices of all values in an array

Viewed 62

I have a matrix A. I would like to generate the indices of all the values in this matrix.

A=np.array([[1,2,3],[4,5,6],[7,8,9]])

The desired output should look like:

[(0,0),(0,1),(0,2),(1,0),(1,1),(2,1),(2,0),(2,1),(2,2)]
5 Answers

You can use:

from itertools import product
list(product(*map(range, A.shape)))

This outputs:

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

Explanation:

A.shape gives the dimensions of the array. For each dimension, we create a range() that generates all of the numbers between 0 and the length of a given dimension. We use map() to perform this for each dimension of the array. Finally, we unpack all of these ranges into the arguments of itertools.product() to create the Cartesian product among all these ranges.

Notably, the use of list unpacking and map() means that this approach can handle ndarrays with an arbitrary number of dimensions. At the time of posting this answer, all of the other answers cannot be immediately extended to a non-2D array.

This should work.

indices = []
for i in range(len(A)):
    for j in range(len(A[i])):
        indices.append((i,j))

Heres a way of doing by using itertools combinations

from itertools import combinations
sorted(set(combinations(tuple(range(A.shape[0])) * 2, 2)))

combinations chooses two elements from the list and pairs them, which results in duplication, so converting it to set to remove duplications and then sorting it.

This line of list comprehension works. It probably isn't as fast as using itertools, but it does work.

[(i,j) for i in range(len(A)) for j in range(len(A[i]))]

Using numpy only you can take advantage of ndindex

list(np.ndindex(A.shape))

or unravel_index:

list(zip(*np.unravel_index(np.arange(A.size), A.shape)))

Output:

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

NB. The second option enables you to pass a order='C' (row-major) or order='F' (column-major) parameter to get a different order of the coordinates

Example on A = np.array([[1,2,3],[4,5,6]])

order='C' (default):

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

order='F':

[(0, 0), (1, 0), (0, 1), (1, 1), (0, 2), (1, 2)]
Related