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

问题

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

回答1:

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]])

来源：https://stackoverflow.com/questions/51243609/slicing-array-by-using-another-array-as-the-slice-indices-along-axis