Reputation: 5696
I have the following array of shape(5,2,3)
, which is a collection of 2 * 3
arrays.
a = array([[[ 0, 2, 0],
[ 3, 1, 1]],
[[ 1, 1, 0],
[ 2, 2, 1]],
[[ 0, 1, 0],
[ 3, 2, 1]],
[[-1, 2, 0],
[ 4, 1, 1]],
[[ 1, 0, 0],
[ 2, 3, 1]]])
1) How can I check if there exists a 2 * 3
array in this array of arrays where at least one element is negative in it?
#which is this:
[[-1, 2, 0],
[ 4, 1, 1]]
2) After that how can I remove the above found 2 * 3
array from a
?
A vectorized implementation is much appreciated but looping is fine too.
Upvotes: 2
Views: 704
Reputation: 221624
You could do -
a[~(a<0).any(axis=(1,2))]
Or the equivalent with .all()
and thus avoid the inverting
-
a[(a>=0).all(axis=(1,2))]
Sample run -
In [35]: a
Out[35]:
array([[[ 0, 2, 0],
[ 3, 1, 1]],
[[ 1, 1, 0],
[ 2, 2, 1]],
[[ 0, 1, 0],
[ 3, 2, 1]],
[[-1, 2, 0],
[ 4, 1, 1]],
[[ 1, 0, 0],
[ 2, 3, 1]]])
In [36]: a[~(a<0).any(axis=(1,2))]
Out[36]:
array([[[0, 2, 0],
[3, 1, 1]],
[[1, 1, 0],
[2, 2, 1]],
[[0, 1, 0],
[3, 2, 1]],
[[1, 0, 0],
[2, 3, 1]]])
Upvotes: 1
Reputation: 107347
Use any
:
In [10]: np.any(a<0,axis=-1)
Out[10]:
array([[False, False],
[False, False],
[False, False],
[ True, False],
[False, False]], dtype=bool)
Or more complete, if you want the corresponding index for (2,3) array:
In [22]: np.where(np.any(a<0,axis=-1).any(axis=-1))
Out[22]: (array([3]),)
# Or as mentioned in comment you can pass a tuple to `any` np.where(np.any(a<0,axis=(1, 2)))
You can also get the array with a simple indexing:
In [27]: a[np.any(a<0, axis=(1, 2))]
Out[27]:
array([[[-1, 2, 0],
[ 4, 1, 1]]])
Upvotes: 1