如何将NumPy二维数组每行指定索引的元素设为指定值？

可以利用NumPy的整数数组索引结合广播特性实现需求，无需循环，效率更高：

代码实现 #

import numpy as np

# 初始化示例数组
arr = np.array([[0, 1, 2],
                [3, 4, 5],
                [6, 7, 8],
                [9, 10, 11]])

indices = np.array([[0, 1],
                    [1, 2],
                    [0, 2],
                    [0, 2]])

# 生成与indices形状匹配的行索引数组
row_indices = np.arange(arr.shape[0])[:, np.newaxis]

# 批量定位并修改元素
arr[row_indices, indices] = -1

print(arr)

原理说明 #

• np.arange(arr.shape[0])生成行号数组[0,1,2,3]，通过[:, np.newaxis]将其转为列向量（形状(4,1)）；
• 列向量形式的行索引会与indices（形状(4,2)）自动广播，形成完整的(行号, 列号)索引对；
• 直接对这些索引位置赋值-1，即可完成按行批量修改元素的操作。

扩展说明 #

如果你的indices是每行长度不一致的不规则数组（比如列表的列表），可以改用循环逐行处理：

for i, cols in enumerate(indices):
    arr[i, cols] = -1

这种写法更直观，适合小数据量场景。

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