Use 2D matrix as indexes for a 3D matrix in numpy?

user_2352847 0
Viewed 3208

Say I have an array of shape 2x3x3, which is a 3D matrix. I also have a 2D matrix of shape 3x3 that I would like to use as indices for the 3D matrix along the first axis. Example is below.

Example run:

>>> np.random.randint(0,2,(3,3)) # index
array([[0, 1, 0],
       [1, 0, 1],
       [1, 0, 0]])

>> np.random.randint(0,9,(2,3,3)) # 3D matrix
array([[[4, 4, 5],
        [2, 6, 7],
        [2, 6, 2]],

       [[4, 0, 0],
        [2, 7, 4],
        [4, 4, 0]]])
>>> np.array([[4,0,5],[2,6,4],[4,6,2]]) # result
array([[4, 0, 5],
       [2, 6, 4],
       [4, 6, 2]])
matrix numpy python
3 Answers

There is a numpy function off-the-shelf: np.choose. It also comes with some handy broadcast options. 

import numpy as np    
cube = np.arange(18).reshape((2,3,3))
sel = np.array([[1, 0, 1], [0, 1, 1], [0,1,0]])

the_selection = np.choose(sel, cube)


>>>the_selection
array([[ 9,  1, 11],
       [ 3, 13, 14],
       [ 6, 16,  8]])

This method works with any 3D array.

user_45350200
Related