akilat90
akilat90

Reputation: 5696

Element wise comparison in an array of arrays in Numpy

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

Answers (2)

Divakar
Divakar

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

Kasravnd
Kasravnd

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

Related Questions