在NumPy中如何无需循环用二维数组按行索引另一二维数组并赋值

直接使用NumPy高级广播索引即可实现，无需任何显式循环： #

实现代码 #

import numpy as np

# 定义存储索引的数组
idx_arr = np.array([[1, 3], [2, 5]])
# 定义初始全零数组
zero_arr = np.zeros((2, 6), dtype=int)

# 核心赋值逻辑
zero_arr[np.arange(idx_arr.shape[0])[:, np.newaxis], idx_arr] = 1

原理解释 #

• np.arange(idx_arr.shape[0])生成和全零数组行数一致的行坐标序列[0, 1]
• 通过[:, np.newaxis]将行坐标序列扩展为形状(2, 1)的二维数组[[0], [1]]，和形状(2, 2)的索引数组广播匹配
• 批量定位到每一行对应的索引位置，一次性完成赋值操作

该方案为纯向量化操作，无任何Python层循环，处理大规模数组时性能远高于循环实现。 #

执行后打印zero_arr即可得到目标结果：

array([[0, 1, 0, 1, 0, 0],
       [0, 0, 1, 0, 0, 1]])

如果你的初始全零数组为浮点类型，赋值后对应位置会变为1.，不影响逻辑，若需要整数类型可在定义数组时指定dtype=int即可。

内容的提问来源于stack exchange，提问作者x1netta