如何利用多数组与掩码高效填充分类器预测结果数组？

需求说明 #

在分类器预测任务中，需要创建一个新数组，规则如下：

• 当mask_pred为False时，填充指定的pred_value
• 当mask_pred为True时：
  • 若mask_z为True且mask_x为False，取x对应位置的值
  • 若mask_z为False且mask_x为True，取z对应位置的值
  • 若mask_z和mask_x均为True，填充np.NaN

示例输入：

import numpy as np

z = np.array([2, 4, 5, 7])
x = np.array([3, 6, 9, 8])
pred_value = 11

# 修正笔误：np.arary → np.array
mask_z = z % 2 == 0  # array([True, True, False, False])
mask_x = x % 2 == 0  # array([False, True, False, True])
mask_pred = np.array([True, True, False, True])

期望输出：array([3, nan, 11, 7])

掩码实现方案 #

利用NumPy的布尔掩码进行向量化操作，避免循环，适合处理大规模数组：

import numpy as np

# 初始化结果数组，先默认填充pred_value，指定float类型以支持NaN存储
y = np.full_like(z, pred_value, dtype=np.float64)

# 定位mask_pred为True的区域
pred_true_mask = mask_pred

# 条件1：mask_pred=True，mask_z=True且mask_x=False → 取x的值
cond1 = pred_true_mask & mask_z & ~mask_x
y[cond1] = x[cond1]

# 条件2：mask_pred=True，mask_z=False且mask_x=True → 取z的值
cond2 = pred_true_mask & ~mask_z & mask_x
y[cond2] = z[cond2]

# 条件3：mask_pred=True，mask_z和mask_x均为True → 填充NaN
cond3 = pred_true_mask & mask_z & mask_x
y[cond3] = np.nan

print(y)
# 输出：array([ 3., nan, 11.,  7.])

逻辑解释 #

1. 初始化数组：用np.full_like创建和目标数组形状一致的数组，默认填充pred_value，指定dtype为float64是因为NaN仅能存储在浮点型数组中。
2. 分条件批量赋值：通过布尔掩码的逻辑组合（&表示逻辑与，~表示逻辑非），精准定位每个规则对应的位置，直接批量赋值，完全避免循环，处理效率远高于循环遍历。

内容的提问来源于stack exchange，提问作者M. Merida-Floriano