NumPy无循环实现二维整数数组按行映射一维数组值的方法

问题复现 #

首先是示例中的数组定义：

import numpy as np

a = 0.001 * np.random.randint(0, 1000, size=(5))
# 示例a取值：array([0.524, 0.311, 0.603, 0.957, 0.923])

b = np.random.randint(0, 2, size=(3, 5))
# 示例b取值：
# array([[1, 1, 1, 0, 1],
#        [0, 0, 1, 0, 0],
#        [0, 0, 1, 0, 1]])

映射规则：
对b中每一行的第i个元素，值为1时替换为a[i]，值为0时替换为1 - a[i]。
以第二行b[1] = [0, 0, 1, 0, 0]为例，映射计算逻辑为：

[1 - 0.524, 1 - 0.311, 0.603, 1 - 0.957, 1 - 0.923]

计算后预期结果：

[0.476, 0.689, 0.603, 0.043, 0.077]

要求不使用for循环，基于NumPy原生能力完成全矩阵运算。 #

实现方法 #

直接利用NumPy广播机制即可实现，两种常用写法都没有显式循环，执行效率为底层C级别：

• 算术运算写法（速度最快）

res = b * a + (1 - b) * (1 - a)

• np.where写法（逻辑更直观）

res = np.where(b, a, 1 - a)

逻辑说明 #

两种写法的核心都是靠NumPy自动广播适配：形状为(5,)的a和1-a会自动扩展为和b（形状(3,5)）同尺寸的数组，逐位置匹配运算：

• 当b某位置值为1：算术写法中1-b项为0，结果等于a[i]；np.where会直接取对应位置的a[i]
• 当b某位置值为0：算术写法中b项为0，结果等于1 - a[i]；np.where会直接取对应位置的1 - a[i]

结果验证 #

用示例值运算后得到的结果为：

array([[0.524, 0.311, 0.603, 0.043, 0.923],
       [0.476, 0.689, 0.603, 0.043, 0.077],
       [0.476, 0.689, 0.603, 0.043, 0.923]])

第二行结果和预期完全一致，运算正确。 #

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