如何优雅高效初始化指定行赋值的三维NumPy零数组？

你可以利用NumPy的高级索引特性，用一行代码完成批量赋值，无需循环遍历，既优雅又高效：

import numpy as np

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

# 核心单行赋值操作
empty[np.arange(empty.shape[0]), row_indexes] = replacement

原理说明 #

• np.arange(empty.shape[0])生成第一个维度的索引序列[0,1,2]，对应三维数组中每个二维子数组的位置；
• 将这个序列与row_indexes配对后，会精准定位到每个二维子数组中需要赋值的行（即(0,1)、(1,1)、(2,0)这三个行位置）；
• 直接将replacement的对应元素批量赋值到这些位置，最终得到你想要的结果：

[[[0 0]
  [1 2]
  [0 0]]
 [[0 0]
  [3 5]
  [0 0]]
 [[5 6]
  [0 0]
  [0 0]]]

内容的提问来源于stack exchange，提问作者Simon P.