如何用向量替代Numpy中的循环实现指定赋值操作？

问题场景 #

原循环代码如下，需要将其替换为无循环的向量式实现：

for i in range(4):
    Q[i, x[i] : x[i] + 2] = 1 / 2

其中Q = np.zeros([4,4])，x = [0,1,0,1]。

你尝试用np.s_构建切片，但不知道如何根据x的随机值自动生成对应切片集合，这里提供可行的无循环方案：

解决方案 #

不需要依赖np.s_，直接通过生成对应的行、列索引来实现批量赋值：

import numpy as np

Q = np.zeros([4,4])
x = np.array([0,1,0,1])

# 生成行索引：每个行索引重复2次，对应每个行要赋值的两个列位置
rows = np.repeat(np.arange(4), 2)
# 生成列索引：每个x元素对应两个连续列，通过广播展开
cols = x[:, np.newaxis] + np.arange(2)
cols = cols.ravel()

# 批量赋值
Q[rows, cols] = 1/2

效果验证 #

执行后Q的结果和原循环完全一致：

array([[0.5, 0.5, 0. , 0. ],
       [0. , 0.5, 0.5, 0. ],
       [0.5, 0.5, 0. , 0. ],
       [0. , 0.5, 0.5, 0. ]])

原理说明 #

• np.repeat(np.arange(4), 2)生成[0,0,1,1,2,2,3,3]，对应每个行需要赋值两次
• x[:, np.newaxis] + np.arange(2)将x转换为列向量后和[0,1]广播相加，得到每个行对应的两个列索引，再用ravel()展平为一维数组，和行索引一一对应
• 利用Numpy的花式索引直接批量定位到需要赋值的位置，完成无循环赋值

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