np.take from 3D matrix given indices of second dimension-CodePudding

given a 3D array:

a = np.arange(3*4*5).reshape(3,4,5)
array([[[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19]],

       [[20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39]],

       [[40, 41, 42, 43, 44],
        [45, 46, 47, 48, 49],
        [50, 51, 52, 53, 54],
        [55, 56, 57, 58, 59]]])

I would like to create the following matrix:

result = 
array([[20, 21, 22, 23, 24],
       [ 5,  6,  7,  8,  9],
       [50, 51, 52, 53, 54],
       [55, 56, 57, 58, 59]])

Using the indices idx = [1,0,2,2] I.e, I would like to "take" per matrix, the row specified in idx, where len(idx)==a.shape[1] and np.max(idx)<a.shape[0] as idx choose from dimension 1.

CodePudding user response：

Given that your array has three dimensions (x,y,z), since you want to take one value for each row in the yth direction, you can do this:

a[idx, range(a.shape[1])]

Output:

array([[20, 21, 22, 23, 24],
       [ 5,  6,  7,  8,  9],
       [50, 51, 52, 53, 54],
       [55, 56, 57, 58, 59]])