如何无需循环基于数组值创建动态掩码？

可以利用numpy的广播机制实现无循环的掩码创建，效率远高于循环方式，核心是通过坐标网格和广播运算定位需要置0的区域。

实现代码 #

import numpy as np

max_value = 5
num = 10
values = np.random.randint(max_value, size=(num))
mask = np.ones((num, max_value, max_value))

# 生成行、列坐标网格
rows, cols = np.ogrid[:max_value, :max_value]
# 通过广播计算每个样本对应的右下角置0区域
row_condition = rows >= (max_value - values[:, None, None])
col_condition = cols >= (max_value - values[:, None, None])
# 对符合条件的位置赋值为0
mask[row_condition & col_condition] = 0

关键逻辑说明 #

• np.ogrid生成的坐标网格能快速获取行、列的索引信息，避免手动构建复杂索引
• 将values扩展为(num,1,1)的形状后，可与rows、cols进行广播运算，一次性计算出所有样本的置0区域掩码
• 最后通过逻辑与操作定位到需要置0的交集区域，直接批量赋值

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