如何在NumPy中从二维数组每行的掩码有效元素里随机选取？

问题描述 #

我有两个形状均为(M,N)的二维数组data和布尔数组mask，需要在data的每一行中随机选一个元素，但该元素在mask对应的位置必须为True。想知道有没有不用逐行循环的实现方式？已知每行中mask为True的元素至少有2个。

期望用法示例（无法运行） #

import numpy as np

data = np.arange(8).reshape((2,4))
mask = np.array([[True, True, True, True], [True, True, False, False]])
# 仅作为期望的用法参考，实际无法运行
selected_data = np.random.choice(data, mask, num_elements=1, axis=1)

有效结果示例 #

selected_data = [0,4]
selected_data = [1,5]
selected_data = [2,5]
selected_data = [3,4]

解决方案 #

这里提供两种纯numpy向量化实现，完全不用逐行循环：

方法一：基于有效索引随机选取 #

import numpy as np

data = np.arange(8).reshape((2,4))
mask = np.array([[True, True, True, True], [True, True, False, False]])

# 提取每行中mask为True的列索引，整理成(M, K)的数组（K是每行有效元素数）
valid_cols = np.where(mask)[1].reshape(mask.shape[0], -1)
# 为每行随机生成一个有效列的索引
rand_idx = np.random.randint(0, valid_cols.shape[1], size=mask.shape[0])
# 按行提取对应元素
selected_data = data[np.arange(mask.shape[0]), valid_cols[np.arange(mask.shape[0]), rand_idx]]

print(selected_data)

方法二：基于累积和定位目标元素 #

import numpy as np

data = np.arange(8).reshape((2,4))
mask = np.array([[True, True, True, True], [True, True, False, False]])

# 计算每行mask的累积和，用于标记每个True是该行的第几个有效元素
mask_cumsum = mask.cumsum(axis=1)
# 为每行随机选一个有效元素的序号（从1开始）
target_k = np.random.randint(1, mask.sum(axis=1)+1)
# 找到每行中第k个True的列位置
selected_cols = np.argmax(mask_cumsum == target_k[:, np.newaxis], axis=1)
# 提取对应元素
selected_data = data[np.arange(data.shape[0]), selected_cols]

print(selected_data)

这两种方法都利用numpy的向量化操作，避免了显式循环，在数据量较大时效率会比循环高很多。

内容的提问来源于stack exchange，提问作者Nagabhushan S N