Question

I have an array of N-dimensional values arranged in a 2D array. Something like:

import numpy as np
data = np.array([[[1,2],[3,4]],[[5,6],[1,2]]])

I also have a single value x that I want to compare against each data point, and I want to get a 2D array of boolean values showing whether my data is equal to x.

x = np.array([1,2])

If I do:

data == x

I get

# array([[[ True,  True],
#        [False, False]],
#
#       [[False, False],
#        [ True,  True]]], dtype=bool)

I could easily combine these to get the result I want. However, I don’t want to iterate over each of these slices, especially when data.shape[2] is larger. What I am looking for is a direct way of getting:

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

Any ideas for this seemingly easy task?

Answer 1

Well, (data == x).all(axis=-1) gives you what you want. It’s still constructing a 3-d array of results and iterating over it, but at least that iteration isn’t at Python-level, so it should be reasonably fast.