如何高效基于条件从两个NumPy数组中按行选取构造新数组？

问题复现 #

原有循环实现逻辑为逐行判断：当a[i] >= b[i]时取y的第i行，否则取x的第i行，直接使用np.select会抛出形状不匹配错误：
ValueError: shape mismatch: objects cannot be broadcast to a single shape. Mismatch is between arg 0 with shape (5,) and arg 1 with shape (5, 4).

报错原因 #

NumPy广播规则要求数组维度从尾部开始对齐：

• 判断条件a >= b是形状为(5,)的一维布尔数组
• 待选取的x、y是形状为(5,4)的二维数组
  一维条件缺少第二个维度，无法和二维数组完成广播匹配，因此触发报错。np.select更适合多条件下逐元素选取的场景，逐行选整行的需求用更简单的API即可实现。

最优实现 #

方法1：np.where + 维度扩展（一行代码实现） #

给一维条件数组新增长度为1的轴，将形状转换为(5,1)，即可沿列方向广播条件，匹配二维数组形状：

import numpy as np

a = np.array([4, 0, 1, 2, 4])
b = np.array([0, 4, 3, 2, 0])
y = np.array([[0, 0, 0, 0],
              [0, 0, 0, 1],
              [0, 0, 1, 0],
              [0, 0, 1, 1],
              [0, 0, 1, 0]])
x = np.array([[0, 0, 0, 0],
              [1, 1, 1, 0],
              [1, 1, 0, 0],
              [1, 1, 1, 1],
              [0, 0, 1, 0]])

z = np.where((a >= b)[:, None], y, x)
print(z)

运行结果和原生循环完全一致，运行效率比Python循环高1~2个数量级。

方法2：布尔索引赋值（逻辑更直观） #

如果觉得维度扩展写法不易理解，可以用布尔索引分步完成赋值，效率和方法1相当：

# 先初始化z为x的副本
z = x.copy()
# 生成行筛选掩码
mask = a >= b
# 将满足条件的行替换为y对应行
z[mask] = y[mask]

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