Numpy 3D array data slicing along with specified axis

Question

Suppose I have an array with shape (3, 4, 5) and want to slice along the second axis with an index array [2, 1, 0].

I could not explain what I want to do in text, so please refer the below code and figure:

>>> src = np.arange(3*4*5).reshape(3,4,5)
>>> index = [2,1,0]

>>> src
>>> 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]]])
>>> # what I need is:
    array([[[10, 11, 12, 13, 14]],  # slice the 2nd row (index[0])
           [[25, 26, 27, 28, 29]],  # 1st row (index[1])
           [[40, 41, 42, 43, 44]]])  # 0th row (index[2])

cs95 · Accepted Answer · 2018-12-26 19:09:51Z

1

src[np.arange(src.shape[0]), [2, 1, 0]]
# src[np.arange(src.shape[0]), [2, 1, 0], :]
array([[10, 11, 12, 13, 14],
       [25, 26, 27, 28, 29],
       [40, 41, 42, 43, 44]])

We need to compute the indices for axis=0:

>>> np.arange(src.shape[0])
array([0, 1, 2])

And we already have the indices for axes=1. We then slice across axis=3 to extract our cross-section.

edited Dec 26, 2018 at 19:09

answered Dec 26, 2018 at 19:01

cs95

406k106 gold badges744 silver badges797 bronze badges

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

Comments

Dani Mesejo · Accepted Answer · 2018-12-26 19:04:17Z

1

You could do:

import numpy as np

arr = np.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]]])


first, second = zip(*enumerate([2, 1, 0]))

result = arr[first, second, :]
print(result)

Output

[[10 11 12 13 14]
 [25 26 27 28 29]
 [40 41 42 43 44]]

answered Dec 26, 2018 at 19:04

Dani Mesejo

62.2k6 gold badges56 silver badges86 bronze badges

Collectives™ on Stack Overflow

Numpy 3D array data slicing along with specified axis

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related