能否优化Numpy矩阵指定列起始赋值实现？寻求向量化方案

已知列索引数组 y = np.array([3, 0, 4, 1])，零矩阵 x = np.zeros(shape=(4, 5))，需要将每行从y对应元素指定的列开始至末尾的元素设为1，最终目标矩阵如下：

[[0. 0. 0. 1. 1.]      
 [1. 1. 1. 1. 1.]  
 [0. 0. 0. 0. 1.]  
 [0. 1. 1. 1. 1.]]

例如，y[0] = 3，则第0行的第3、4列需设为1。目前已通过循环实现：

for (idx, num) in enumerate(y): 
    x[idx, num:] = 1

请问能否使用Numpy向量化方法改写或优化该实现？

几种向量化实现方案 #

方案1：广播生成布尔掩码（最推荐） #

通过生成列索引数组，和y进行广播对比，直接得到需要赋值的布尔掩码，无需显式循环：

import numpy as np

y = np.array([3, 0, 4, 1])
x = np.zeros(shape=(4, 5))

# 生成列索引，广播后和y的每个元素比较
cols = np.arange(x.shape[1])
mask = cols >= y[:, np.newaxis]  # 形状(4,5)的布尔矩阵
x[mask] = 1

这种方法完全向量化，矩阵规模越大，对比循环的效率优势越明显，且代码可读性强。

方案2：用np.where提取目标索引 #

基于和方案1相同的掩码逻辑，通过np.where提取出所有需要设为1的行、列索引，再批量赋值：

import numpy as np

y = np.array([3, 0, 4, 1])
x = np.zeros(shape=(4, 5))

cols = np.arange(x.shape[1])
row_idx, col_idx = np.where(cols >= y[:, np.newaxis])
x[row_idx, col_idx] = 1

效果和方案1一致，适合需要单独获取目标索引的场景。

方案3：累积求和构造连续1 #

通过标记每行的起始位置，再利用累积求和生成连续的1序列：

import numpy as np

y = np.array([3, 0, 4, 1])
x = np.zeros(shape=(4, 5))

# 先在每行的起始列标记1
rows = np.arange(x.shape[0])
x[rows, y] = 1
# 反向累积求和再反转，得到从起始列到末尾的连续1
x = x.cumsum(axis=1)[:, ::-1].cumsum(axis=1)[:, ::-1]

这种方法思路巧妙，但可读性稍弱，性能和前两种方案接近。

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