Numpy如何基于布尔掩码对3D矩阵沿axis0求条件均值
2026-8-27
问题描述
我有一个如下所示的3D Numpy矩阵A：
import numpy as np A = np.array([ [ [90, 84, 88], [10, 30, 17], [7, 0, 4] ], [ [88, 83, 102], [12, 14, 15], [12, 17, 7] ], [ [150, 90, 103], [9, 16, 21], [17, 7, 12] ] ]) A
输出为：
array([[[ 90, 84, 88], [ 10, 30, 17], [ 7, 0, 4]], [[ 88, 83, 102], [ 12, 14, 15], [ 12, 17, 7]], [[150, 90, 103], [ 9, 16, 21], [ 17, 7, 12]]])
我沿axis 0对矩阵A的元素做差，生成差值矩阵B：
B = A.T[..., None, :] - A.T[...,None] B
输出为：
array([[[[ 0, -2, 60], [ 2, 0, 62], [-60, -62, 0]], [[ 0, 2, -1], [ -2, 0, -3], [ 1, 3, 0]], [[ 0, 5, 10], [ -5, 0, 5], [-10, -5, 0]]], [[[ 0, -1, 6], [ 1, 0, 7], [ -6, -7, 0]], [[ 0, -16, -14], [ 16, 0, 2], [ 14, -2, 0]], [[ 0, 17, 7], [-17, 0, -10], [ -7, 10, 0]]], [[[ 0, 14, 15], [-14, 0, 1], [-15, -1, 0]], [[ 0, -2, 4], [ 2, 0, 6], [ -4, -6, 0]], [[ 0, 3, 8], [ -3, 0, 5], [ -8, -5, 0]]]])
随后我移除B的上三角元素并取绝对值，将得到的新矩阵存储为C：
C = np.tril(B, k=0) C = np.absolute(C) C
输出为：
array([[[[ 0, 0, 0], [ 2, 0, 0], [60, 62, 0]], [[ 0, 0, 0], [ 2, 0, 0], [ 1, 3, 0]], [[ 0, 0, 0], [ 5, 0, 0], [10, 5, 0]]], [[[ 0, 0, 0], [ 1, 0, 0], [ 6, 7, 0]], [[ 0, 0, 0], [16, 0, 0], [14, 2, 0]], [[ 0, 0, 0], [17, 0, 0], [ 7, 10, 0]]], [[[ 0, 0, 0], [14, 0, 0], [15, 1, 0]], [[ 0, 0, 0], [ 2, 0, 0], [ 4, 6, 0]], [[ 0, 0, 0], [ 3, 0, 0], [ 8, 5, 0]]]])
我的目标是使用np.mean(A, axis=0)沿axis 0对A求均值，但仅当C中对应位置的取值处于名为limit的阈值范围内时，对应元素才参与计算。
我已构建如下掩码：
limit = 8 E = np.where(C <= limit, True, False) E
输出为：
array([[[[ True, True, True], [ True, True, True], [False, False, True]], [[ True, True, True], [ True, True, True], [ True, True, True]], [[ True, True, True], [ True, True, True], [False, True, True]]], [[[ True, True, True], [ True, True, True], [ True, True, True]], [[ True, True, True], [False, True, True], [False, True, True]], [[ True, True, True], [False, True, True], [ True, False, True]]], [[[ True, True, True], [False, True, True], [False, True, True]], [[ True, True, True], [ True, True, True], [ True, True, True]], [[ True, True, True], [ True, True, True], [ True, True, True]]]])
请问如何仅基于E中的True值，沿axis 0对A中对应元素求均值？所有False对应的元素不参与均值计算，仅True关联的元素纳入计算范围。如果存在更简洁的实现方式，也欢迎告知。
解决方案
首先明确维度对应关系：A形状为(3,3,3)，轴顺序是(axis0, 行, 列)；掩码E是对转置后的A做两两差生成的4维矩阵，轴顺序为(列, 行, 轴0索引1, 轴0索引2)，最后两个维度对应axis0上样本的两两配对。
计算逻辑为：对每个固定的(行,列)位置，仅保留和其他所有样本差值都在阈值范围内的axis0样本，再对这些有效样本求均值。
方法1：基于已生成的掩码E计算
先把4维的配对掩码归约为和A同形状的有效性掩码，再计算加权平均：
# 对两两配对的掩码做逻辑与，标记每个元素是否和其他所有元素差值都符合阈值 valid_mask = E.all(axis=(-1, -2)) # 转置回和A一致的轴顺序，形状与A完全相同 valid_mask = valid_mask.transpose(2, 1, 0) # 计算过滤后的均值：有效元素求和 / 有效元素个数 filtered_sum = np.where(valid_mask, A, 0).sum(axis=0) valid_count = valid_mask.sum(axis=0) mean_result = filtered_sum / valid_count
运行得到结果：
array([[89. , 83.5 , 95. ], [10.33333333, 20. , 17.66666667], [12. , 8. , 7.66666667]])
方法2：更简洁的无中间变量实现
可以跳过构造B、C、E的步骤，直接基于原始矩阵计算，内存占用更低：
limit = 8 # 直接计算axis0上样本的两两差值，生成有效掩码 pair_diff = np.abs(A[:, None, ...] - A[None, ...]) valid_mask = (pair_diff <= limit).all(axis=0) # 计算均值 mean_result = (A * valid_mask).sum(axis=0) / valid_mask.sum(axis=0)
计算结果和方法1完全一致。如果存在某位置所有样本都被判定为无效的情况，可以用np.divide的where参数避免除0警告，给无效位置填充指定默认值（比如NaN）。
内容的提问来源于stack exchange，提问作者equanimity
相关产品推荐
相关产品推荐
