如何向量化实现Numpy按元素概率二选一操作，替代循环提升性能

向量化实现方案 #

你可以借助numpy的内置随机生成和条件选择函数实现完全无循环的计算，核心逻辑如下：

1. 生成与P形状完全相同、取值在[0,1)区间的均匀随机数组rand_mask
2. 利用布尔判断生成掩码：当rand_mask[i,j] <= P[i,j]时选取A[i,j]，否则选取B[i,j]。由于均匀随机数小于等于P[i,j]的概率恰好等于P[i,j]，完全匹配需求。

实现代码如下：

import numpy as np
rows = 4
columns = 5
A = np.zeros((rows, columns)) # 全0数组
B = np.ones_like(A)           # 全1数组
P = np.random.rand(*A.shape)  # [0,1)区间随机概率数组

# 无循环向量化实现
rand_mask = np.random.rand(*P.shape)
C = np.where(rand_mask <= P, A, B)

该方案所有运算都在numpy底层C层面执行，处理大数组时性能相比双层Python循环提升可达上千倍。

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