CODEWITHSUNDEEP
pythonnumpy
HappyPy
HappyPy

Reputation: 10697

get indices of numpy multidimensional arrays

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

I'm trying to get the indices of arr when arr==1.

I thought this would work but it doesn't give the expected output:

>>> np.where(arr==1)
   (array([0], dtype=int64), array([0], dtype=int64), array([0], dtype=int64))

Upvotes: 0

Views: 259

Answers (3)

Valdi_Bo
Valdi_Bo

Reputation: 30971

An interesting option is to use argwhere to get indices of elements in question.

To present a more instructive example, assume that we want indices of your arr where the value is <= 3.

The code to get them is np.argwhere(arr <= 3), getting:

array([[0, 0, 0],
       [0, 0, 1],
       [0, 1, 0]], dtype=int64)

​Meaning that the "wanted" elemens are:

arr[0, 0, 0], arr[0, 0, 1], arr[0, 1, 0],

If you want e.g. to print indices of the "wanted" elements and their values, you can run:

for ind in np.argwhere(arr <= 3):
    print(f'{ind}: {arr.__getitem__(tuple(ind))}')

The result is:

[0 0 0]: 1
[0 0 1]: 2
[0 1 0]: 3

Note also that my code works regardless of the number of dimensions in your array, whereas in other solutions the number of dimensions is somehow "fixed" in the code.

Upvotes: 0

bousof
bousof

Reputation: 1251

For multiple numbers, you can mix np.where and np.isin functions:

Here is an example:

import numpy as np

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

loc_i, loc_j, loc_k = np.where(np.isin(arr, val))
print('locations:')
[print(f'({loc_i[i]},{loc_j[i]},{loc_k[i]})') for i in range(loc_i.size)]

Upvotes: 1

dimay
dimay

Reputation: 2804

If you change arr:

arr = np.array([[[1,2],[3,4]],[[1,1],[7,8]]])

and you will get

np.where(arr==1)
# (array([0, 1, 1]), array([0, 0, 0]), array([0, 0, 1]))

it is mean:

arr[0][0][0] == 1
arr[1][0][0] == 1
arr[1][0][1] == 1

When you want display coordinate in one row:

np.array(np.where(arr==1)).T
# array([[0, 0, 0],[1, 0, 0],[1, 0, 1]])

Upvotes: 1

Related Questions