Setting Values in an Multidimensional numpy Array with Indexes-CodePudding

i am searching for an elegant and efficient way to initialize a 3dimensional numpy array that is all zeros but for a specific row. E.g.

empty = np.zeros((3,3,2))
row_indexes = np.array([1,1,0])
replacement = np.array([[1,2][3,5][5,6]])
# intended outcome:
[[[0 0]
  [1 2]
  [0 0]]
 [[0 0]
  [3 5]
  [0 0]]
 [[5 6]
  [0 0]
  [0 0]]]

I could iterate over each two dimensional array along the zero dimension, but that seems inelegant. Is there an efficient oneliner for this problem?

CodePudding user response：

The row_indexes is used as the second dimension here. Just build the index of the corresponding first dimension, and then directly assign values:

>>> empty[np.arange(len(empty)), row_indexes] = replacement
>>> empty
array([[[0., 0.],
        [1., 2.],
        [0., 0.]],

       [[0., 0.],
        [3., 5.],
        [0., 0.]],

       [[5., 6.],
        [0., 0.],
        [0., 0.]]])