带条件的NumPy广播：筛选满足索引条件的2x2数组

问题分析 #

你已经通过广播得到了形状为(4,4,2,2)的数组c，其中：

• c[i,j]对应b的第i个2x2数组（b的axis=0索引为i）与a的第j个2x2数组（a的axis=1索引为j）相加的结果
• 需要筛选出所有满足j > i的c[i,j]子数组，最终组成形状为(6,2,2)的数组

实现步骤 #

1. 生成索引网格：创建对应c前两个维度的索引矩阵，明确i（b的axis=0索引）和j（a的axis=1索引）的位置
2. 创建布尔掩码：根据j > i的条件生成布尔数组，标记需要保留的子数组位置
3. 筛选并重塑结果：用掩码提取符合条件的子数组，Numpy会自动处理维度扁平化，直接得到目标形状

完整代码 #

import numpy as np

# 原始代码生成a、b、c
a = np.random.randint(0, 10, size=(4, 2, 2))   
b = np.random.randint(0, 10, size=(4, 2, 2))

a = a[np.newaxis, :]  # shape变为(1,4,2,2)
b = b[:, np.newaxis]  # shape变为(4,1,2,2)

c = a + b  # 广播后shape为(4,4,2,2)

# 筛选满足j > i的子数组
i, j = np.indices((4, 4))  # 生成(4,4)的索引矩阵，i对应c的axis=0，j对应c的axis=1
mask = j > i  # 布尔掩码，标记符合条件的位置
result = c[mask]  # 提取符合条件的子数组，自动变为(6,2,2)

# 验证结果形状
print(result.shape)  # 输出 (6, 2, 2)

补充说明 #

• 掩码mask中True的数量可通过mask.sum()查看，结果为6，与题目描述一致
• 若需筛选其他条件，只需修改掩码条件即可：
  • j < i会得到6个符合条件的子数组
  • j == i会得到4个符合条件的子数组

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