Is there a simple numpy way to broadcast one dimension and look-up using others?-CodePudding

I have 2 numpy arrays with dimensions NxM and Nxa. I would like to extract 'a' of the 'M' elements from each row of the NxM matrix with row-indices given by the Nxa. It seems like there should be a simple broadcasting solution but I can't find it.

Example given below is hopefully demonstrative.

import numpy as np
N=5
M=7
a=3

NxM = np.array([[0, 1, 2, 3, 4, 5, 6],
                [10, 11, 12, 13, 14, 15, 16],
                [20, 21, 22, 23, 24, 25, 26],
                [30, 31, 32, 33, 34, 35, 36],
                [40, 41, 42, 43, 44, 45, 46]])

Nxa = np.array([[1, 3, 6],
                [0, 1, 2],
                [1, 4, 5],
                [2, 3, 4],
                [3, 5, 6]])

#desired output
output = [[1, 3, 6],
          [10, 11, 12],
          [21, 24, 25],
          [32, 33, 34],
          [43, 45, 46]]

CodePudding user response：

You are looking for fancy indexing:

NxM[np.arange(N).reshape(-1, 1), Nxa]

The shape of an output array is the shape of the broadcasted indices.

CodePudding user response：

numpy.take_along_axis is what you are looking for.

np.take_along_axis(NxM, Nxa, axis=1)

array([[ 1,  3,  6],
       [10, 11, 12],
       [21, 24, 25],
       [32, 33, 34],
       [43, 45, 46]])