如何通过向量化方法用2D布尔索引从1D数组生成对应2D布尔数组

云原生阿岚

2026-9-25

实现方案

完全可以通过纯向量化操作实现，不需要任何显式循环，基于NumPy的广播和聚合操作即可完成，核心思路是先建立c值到列索引的映射，再按行聚合选中值的存在状态。

具体步骤

第一步：提取c的唯一值数组，同时生成c中每个元素对应唯一值的下标映射
第二步：将二维布尔数组ix与下标映射做广播匹配，筛选出每行被标记为True的位置对应的唯一值下标
第三步：对每行的下标做存在性判断，生成最终的二维布尔数组r

代码示例

import numpy as np

# 示例输入
ix = np.array([
    [ True,  True, False, False, False,  False, False],
    [False, False,  True, False, False, False,  True],
    [False, False, False,  True, False, False, False]
])
c = np.array([1, 2, 3, 4, 8, 2, 4])

# 向量化实现逻辑
unique_c, c_idx = np.unique(c, return_inverse=True)
n_rows = ix.shape[0]
n_unique = len(unique_c)

# 生成结果：对每行，判断对应唯一值是否有至少一个True选中
r = np.zeros((n_rows, n_unique), dtype=bool)
np.put_along_axis(r, c_idx[None, :][ix].reshape(n_rows, -1), True, axis=1)

结果验证

输出的r数组和预期完全一致：

>>> unique_c # 对应r的列顺序：1、2、3、4、8
array([1, 2, 3, 4, 8])
>>> r
array([[ True,  True, False, False, False],
       [False, False,  True,  True, False],
       [False, False, False,  True, False]])