基于数值比较高效查找3D矩阵中符合多参数条件的样本-位置对

问题背景 #

给定3D NumPy矩阵：

import numpy as np
matrix = np.random.randn(100,10,500)

各维度定义：

• X维度：数据样本（共100个）
• Y维度：参数/变量类型（共10个）
• Z维度：位置（共500个）

需求是找出满足以下规则的（样本索引，位置索引）元组数组：

1. 可针对单个/多个参数设置数值范围（如-1.0至5.5）或等值条件（如等于1.39），支持多条件或逻辑组合
2. 支持对条件取反（NOT）
3. 最终结果需同时满足所有参数的各自条件

之前采用循环结合np.where的实现效率极低，以下是高效的向量化解决方案。

核心思路：用NumPy向量化操作替代循环 #

NumPy的底层C实现向量化操作，能彻底避开Python循环的性能开销。我们可以为每个参数生成布尔掩码，再将所有参数的掩码做与运算，最后批量提取满足条件的索引对。

步骤1：编写条件掩码生成函数 #

这个函数可灵活支持范围、等值、取反、或逻辑组合，直接针对指定参数生成2D布尔掩码（样本×位置）：

def create_mask(matrix, param_idx, conditions, negate=False):
    """
    针对指定参数生成布尔掩码
    :param matrix: 3D NumPy矩阵 (样本数, 参数数, 位置数)
    :param param_idx: 目标参数的Y维度索引
    :param conditions: 条件列表，每个元素格式：
                       - 等值: ('eq', 目标值)
                       - 范围比较: ('gt'/'lt'/'ge'/'le', 阈值)
                       - 区间: ('between', 最小值, 最大值)
    :param negate: 是否对结果取反
    :return: 2D布尔掩码 (样本数, 位置数)
    """
    param_data = matrix[:, param_idx, :]
    mask = np.zeros_like(param_data, dtype=bool)
    
    for cond in conditions:
        op = cond[0]
        if op == 'eq':
            cond_mask = param_data == cond[1]
        elif op in ('gt', 'lt', 'ge', 'le'):
            cond_mask = getattr(param_data, op)(cond[1])
        elif op == 'between':
            cond_mask = (param_data >= cond[1]) & (param_data <= cond[2])
        else:
            raise ValueError(f"不支持的操作符: {op}")
        
        mask |= cond_mask  # 或逻辑组合多个条件
    
    if negate:
        mask = ~mask
    
    return mask

步骤2：组合所有参数的条件掩码 #

以示例条件为例：

• 参数0：值介于-1.0到5.5 或 等于1.39
• 参数3：值大于2.0（取反后即≤2.0）
• 参数7：值不等于0.0

通过逐个生成参数掩码并做与运算，得到最终的全局掩码：

# 定义各参数的条件规则
param_conditions = [
    (0, [('between', -1.0, 5.5), ('eq', 1.39)]),
    (3, [('gt', 2.0)], True),
    (7, [('eq', 0.0)], True)
]

# 初始化全局掩码为全True
final_mask = np.ones((matrix.shape[0], matrix.shape[2]), dtype=bool)

# 叠加所有参数的条件
for param_idx, conditions, negate in param_conditions:
    param_mask = create_mask(matrix, param_idx, conditions, negate)
    final_mask &= param_mask  # 同时满足所有参数条件

步骤3：提取目标索引对 #

用np.where批量提取掩码中为True的位置，直接得到样本和位置的索引数组，再组合成元组：

sample_indices, pos_indices = np.where(final_mask)
# 转换为（样本索引，位置索引）元组数组
result = np.array(list(zip(sample_indices, pos_indices)))

方案优势 #

• 全程使用NumPy向量化操作，避免Python循环的性能损耗
• 布尔掩码运算基于底层C实现，处理百万级数据也能快速完成
• np.where批量提取索引，无需逐个遍历判断

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