Slicing array by using another array as the slice indices along axis

Question

Say I have an array that looks like the following:

arr = [[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]]

And I have another array slicer = [1,3,2]. I want to apply these values as the slice index over axis 0 measure along axis 1.

This doesn't work (and in fact contains no way of specifying that the along part is axis 1 in an ndarray) but suppose I tried arr[:slicer, :]

I would hope to obtain,

out = [[1,   2,   3],
       [nan, 5,   6],
       [nan, 8, nan]]

which is the combination of applying the slice arr[:1, :], arr[:3, :], arr[:2, :] and then selecting from those the 1st, 2nd and 3rd columns respectively and reassembling into the array above, dropping missing values.

I want to avoid loops and trying to find a fast vectorised solution

Nils Werner · Accepted Answer · 2018-07-09 11:05:20Z

7

For this operation you need to first generate a boolean index mask that marks all fields you want to set to nan. Broadcasting makes it easy to perform an "outer comparison" that yields the desired result

slicer = numpy.asarray([1, 3, 2])
mask = numpy.arange(3)[:, None] >= slicer
mask
# array([[False, False, False],
#        [ True, False, False],
#        [ True, False,  True]])

You can then simply use this mask to index data

data = numpy.arange(1, 10, dtype=float).reshape(3, 3)
data[mask] = numpy.nan
data
# array([[ 1.,  2.,  3.],
#        [nan,  5.,  6.],
#        [nan,  8., nan]])

answered Jul 9, 2018 at 11:05

Nils Werner

37.2k7 gold badges85 silver badges108 bronze badges

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

1 Comment

Attack68 Over a year ago

marvellous result, in record time. Thanks.

Collectives™ on Stack Overflow

Slicing array by using another array as the slice indices along axis

1 Answer 1

1 Comment

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Linked

Related