H.Rappeport
H.Rappeport

Reputation: 607

Get indices of slices with at least one element obeying some condition

I have an ndarray A of shape (n, a, b) I want a Boolean ndarray X of shape (a, b) where

X[i,j]=any(A[:, i, j] < 0)

How to achieve this?

Upvotes: 1

Views: 21

Answers (1)

mrzo
mrzo

Reputation: 2135

I would use an intermediate matrix and the sum(axis) method:

np.random.seed(24)

# example matrix filled either with 0 or -1:
A = np.random.randint(2, size=(3, 2, 2)) - 1

# condition test:
X_elementwise = A < 0

# Check whether the conditions are fullfilled at least once:
X = X_elementwise.sum(axis=0) >= 1

Values for A and X:

A = array([[[-1,  0],
        [-1,  0]],

       [[ 0,  0],
        [ 0, -1]],

       [[ 0,  0],
        [-1,  0]]])

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

Upvotes: 1

Related Questions