Set 3D numpy array value to 0 if last axis index is smaller than value in another 2D array

Question

I have a 3D array a with shape (m, n, p) and a 2D array idx with shape (m, n). I want all elements in a where the last axis index is smaller than the corresponding element in idx to be set to 0.

The following code works. My question is : is there a more efficient approach?

a = np.array([[[1, 2, 3],
               [4, 5, 6]],

              [[7, 8, 9],
               [10, 11, 12]],

              [[21, 22, 23],
               [25, 26, 27]]])
idx = np.array([[2, 1],
                [0, 1],
                [1, 1]])
for (i, j), val in np.ndenumerate(idx):
    a[i, j, :val] = 0

The result is

array([[[ 0,  0,  3],
        [ 0,  5,  6]],

       [[ 7,  8,  9],
        [ 0, 11, 12]],

       [[ 0, 22, 23],
        [ 0, 26, 27]]])

Divakar · Accepted Answer · 2018-03-01 13:52:28Z

3

Use broadcasting to create the 3D mask and then assign zeros with boolean-indexing -

mask = idx[...,None] > np.arange(a.shape[2])
a[mask] = 0

Alternatively, we can also use NumPy builtin for outer-greater comparison to get that mask -

mask = np.greater.outer(idx, np.arange(a.shape[2]))

Run on given sample -

In [34]: mask = idx[...,None] > np.arange(a.shape[2])

In [35]: a[mask] = 0

In [36]: a
Out[36]: 
array([[[ 0,  0,  3],
        [ 0,  5,  6]],

       [[ 7,  8,  9],
        [ 0, 11, 12]],

       [[ 0, 22, 23],
        [ 0, 26, 27]]])

answered Mar 1, 2018 at 13:52

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Set 3D numpy array value to 0 if last axis index is smaller than value in another 2D array

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related