CODEWITHSUNDEEP
pythonnumpynumpy-ndarray
user6453877
user6453877

Reputation: 334

Remove elements in an ndarray based on condition on one dimension

In a Numpy ndarray, how do I remove elements in a dimension based on condition in a different dimension?

I have:

[[[1 3]
  [1 4]]

 [[2 6]
  [2 8]]

 [[3 5]
  [3 5]]]

I want to remove based on condition x[:,:,1] < 7

Desired output ([:,1,:] removed):

[[[1 3]
  [1 4]]

 [[3 5]
  [3 5]]]

EDIT: fixed typo

Upvotes: 1

Views: 304

Answers (2)

Andy L.
Andy L.

Reputation: 25239

As your desired output, you filter x on axis=0. Therefore, you may try this way

m = (x[:,:,1] < 7).all(1)
x_out = x[m,:,:]

Or simply

x_out = x[m]   

Out[70]:
array([[[1, 3],
        [1, 4]],

       [[3, 5],
        [3, 5]]])

Upvotes: 1

9769953
9769953

Reputation: 12201

This may work:

x[np.where(np.all(x[..., 1] < 7, axis=1)), ...]

yields

array([[[[1, 3],
         [1, 4]],

        [[3, 5],
         [3, 5]]]])

You do get an extra dimension, but that's easy to remove:

np.squeeze(x[np.where(np.all(x[..., 1] < 7, axis=1)), ...])

Briefly how it works:

First the condition: x[..., 1] < 7.
Then test if the condition is valid for all elements along the specific axis: np.all(x[..., 1] < 7, axis=1).
Then, use where to grab the indices instead of an array of booleans: np.where(np.all(x[..., 1] < 7, axis=1)).
And insert those indices into the relevant dimension: x[np.where(np.all(x[..., 1] < 7, axis=1)), ...].

Upvotes: 2

Related Questions