Python中如何基于一维数组条件快速设置二维数组行值（无循环）

要实现你想要的效果，关键是利用Numpy的广播机制，把一维数组b扩展成和二维数组a匹配的形状，这样就能直接进行逐元素的条件判断，完全不需要循环，效率拉满。

直接上可运行的代码：

import numpy as np

a = np.array([[1, 3, 3], [2, 5, 5], [3, 7, 7]])
b = np.array([1, 2, 3])

# 把b转换成列向量，让它能和a的每行广播匹配
b_reshaped = b.reshape(-1, 1)
# 直接进行条件赋值
a[a == b_reshaped] = 0

print(a)

运行后输出就是你想要的：

[[0 3 3]
 [0 5 5]
 [0 7 7]]

为什么这样可行？ #

• 当你把b用reshape(-1, 1)转换成(3,1)形状的列数组后，Numpy会自动把它广播成和a一样的(3,3)形状（相当于把每个元素复制成一行），这样a == b_reshaped就会生成一个同样形状的布尔数组，标记出每行中等于b对应元素的位置。
• 然后用这个布尔数组作为索引，直接把对应位置的值设为0，整个过程都是Numpy底层的C级别的向量化操作，比Python循环快几个数量级，完全适合大规模数组的场景。

另一种等价写法（更简洁） #

你也可以不用显式reshape，直接利用[:, None]来给b增加一个维度，效果是一样的：

a[a == b[:, None]] = 0

这样写更短，而且可读性也不错，推荐在实际代码中使用这种方式。

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