如何基于概率数组生成由0和1组成的NumPy数组？

你可以使用NumPy的np.random.binomial函数实现等效功能，它的向量化操作比random.choices的列表推导高效得多，尤其适合大规模数据场景。

等效代码实现 #

import numpy as np

# 生成概率数组（对应原代码的列表推导）
array_probabilities = np.full(4, 0.5)
print(array_probabilities)  # 输出：[0.5 0.5 0.5 0.5]

# 按概率生成0/1数组
a = np.random.binomial(n=1, p=array_probabilities, size=len(array_probabilities))
print(a)  # 示例输出：[1 0 1 1]

代码说明 #

• np.full(4, 0.5)：快速生成指定长度、元素全为0.5的NumPy数组，替代原代码的列表推导，效率更高。
• np.random.binomial：
  • n=1：表示单次伯努利试验（只有0或1两种结果），正好匹配你需要的0/1生成逻辑；
  • p=array_probabilities：每个位置的概率值，对应生成1的概率（生成0的概率即为1-p）；
  • size：指定输出数组的长度，和概率数组长度保持一致即可。

额外说明 #

如果你的概率数组是普通Python列表，无需提前转成NumPy数组，直接传入p参数即可：

array_probabilities = [0.5 for _ in range(4)]
a = np.random.binomial(1, p=array_probabilities, size=len(array_probabilities))

如果需要将生成的NumPy数组转回Python列表，可使用a.tolist()方法。

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