如何无循环实现用布尔掩码列拆分同形数据数组为多个子数组？

给定形状为(m,n)的数据数组和同形状的布尔掩码数组，能否无需循环完成以下操作：

1. 将掩码数组的每一列应用于数据数组，得到子数组；
2. 存储所有子数组。

循环实现方案 #

import numpy as np

arr = np.array([
    [1,2,3],
    [np.nan,5,6]
])

masks = np.array([
    [True, True, True],
    [False, True, True]
])

arrays = []

for i in range(masks.shape[1]):
    mask_i = masks[:, i]
    arr_i = arr[mask_i, :]
    arrays.append(arr_i)

尝试过的无效方案 #

• np.take：因indices参数要求形状一致，但不同列筛选后选取的行数不同；
• 花式索引arr[masks,:]：不适用于按列处理高维掩码，会返回所有满足True条件的元素而非按列分组的子数组；
• np.split：仅能拆分完整不重叠的子数组，无法根据布尔掩码筛选行。

由于不同列筛选后得到的子数组行数可能不一致，无法用单一numpy数组存储结果，最终只能返回列表。可以用列表推导式实现无显式循环的写法，代码如下：

import numpy as np

arr = np.array([
    [1,2,3],
    [np.nan,5,6]
])

masks = np.array([
    [True, True, True],
    [False, True, True]
])

# 无显式循环的实现
arrays = [arr[masks[:, i], :] for i in range(masks.shape[1])]

说明 #

列表推导式属于Python的隐式循环语法，相比显式for循环更简洁高效，且无需手动编写循环结构，完全满足需求。

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