如何无需遍历列填充Numpy数组中特定区间的值

给定Numpy数组：

import numpy as np

arr = np.array([
    [3., 4., 5., 6., 7., 8., 9., 8., 7., 6.],
    [0., 0., 1., 0., 0., 1., 0., 0., 0., 0.],
    [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]
])

需要将第2行按第1行的1作为分界点填充：

• 从起始到第一个1的位置（含）：填充第0行第一个1对应位置的值
• 第一个1之后到第二个1的位置（含）：填充第0行第二个1对应位置的值
• 第二个1之后到末尾：填充第0行最后一个位置的值

实现代码（无列遍历） #

# 获取第1行中1的索引位置
ones_indices = np.where(arr[1] == 1)[0]

# 生成每个填充区间的起始和结束索引
starts = np.concatenate([[0], ones_indices + 1])
ends = np.concatenate([ones_indices + 1, [arr.shape[1]]])

# 获取每个区间对应的填充值：前两个是第0行对应1的位置的值，最后一个是第0行末尾值
fill_values = np.concatenate([arr[0, ones_indices], [arr[0, -1]]])

# 重复填充值对应区间长度，赋值给第2行
arr[2] = np.repeat(fill_values, ends - starts)

运行后输出目标数组：

print(arr)
# 结果：
# [[3. 4. 5. 6. 7. 8. 9. 8. 7. 6.]
#  [0. 0. 1. 0. 0. 1. 0. 0. 0. 0.]
#  [5. 5. 5. 8. 8. 8. 6. 6. 6. 6.]]

说明 #

全程使用Numpy内置函数实现向量化操作，没有循环遍历列，效率远高于逐列处理：

• np.where快速定位分界点1的位置
• np.concatenate拼接区间的起始/结束索引
• np.repeat根据区间长度批量生成填充值，直接赋值给目标行

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