How to find the index of an element in an numpy array?

Question

I have an array df in which each element is a list of 2 numbers. Given an element p = [18, 169]. I would like to find the indices of such elements p in df. Given df

[[[13, 169],    [18, 169],  [183, 169]],
 [[-183, 169],  [18, 169],  [183, 169]],
 [[18, 169],    [-18, 169], [183, 169]]]

With (df == p).all(-1), I get

array([[False,  True, False],
       [False,  True, False],
       [ True, False, False]])

What I want is

[[0, 1],
 [1, 1],
 [2, 0]]

Could you please elaborate on how to do so?

import numpy as np
df = np.array([[[13, 169],   [18, 169], [183, 169]],
               [[-183, 169], [18, 169], [183, 169]],
               [[18, 169],   [-18, 169], [183, 169]]])
p = [18, 169]
ind = (df == p).all(-1)
ind

df is not a 2D array, but a 3D array. You can check this via: df.ndim, which returns 3. — James
– James, Commented Apr 10, 2021 at 20:57
@James I posted an answer with np.stack(np.where((df == p).all(-1))). It produces what I'm looking for. I wonder why you deleted it. — Akira
– Akira, Commented Apr 10, 2021 at 21:03
I deleted my answer because @Hans Musgrave had a better answer using a shorter version of the numpy API. — James
– James, Commented Apr 10, 2021 at 21:05

Hans Musgrave · Accepted Answer · 2021-04-10 21:00:00Z

4

What you've computed with (df==p).all(-1) is a mask. They have lots of uses, but you can use that directly to compute the value you want.

# True or false at each coordinate
mask = (df==p).all(-1)

# Extract the coordinates where the mask is True
result = np.argwhere(mask)

answered Apr 10, 2021 at 21:00

Hans Musgrave

7,1912 gold badges20 silver badges40 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Akira Over a year ago

So elegant. Thank you so much!!

Collectives™ on Stack Overflow

How to find the index of an element in an numpy array?

1 Answer 1

1 Comment

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Related