使用numpy.where()修改数组时0值未正确更新的问题

问题现象 #

用numpy生成60个0/1的伯努利随机数组正常，但用numpy.where()修改数组时出现两个异常：

• 原数组中1成功替换为"Z"
• 原数组中0本该替换为新的二项分布随机数（0或1），但实际全是0；后续统计新生成的1时结果始终为0

原代码 #

import numpy

n = 50 #运行重复次数
p = 0.1
count = 0

# 模拟第一天
result = numpy.random.binomial(1, p, 60) 
print (result)
for ele in result: 
        if (ele == 1):
            count = count +1
print ("Number Infected:", count)

# 创建追踪感染情况的矩阵
result = numpy.where(result == 1, "Z", numpy.random.binomial(1,p*count))
count = 0
for ele in result: 
        if (ele == 1):
            count = count +1

# while count > 0: 最终会进入该循环条件 #
# 模拟第i天 #
#for ele in result:
#        if (ele == 0):
#            ele == "X"
#        else:
#            ele == "I"
print(result)
print("Number Infected:", count)

当前运行结果 #

[0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0
 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 0 0 0 1 0 0]
Number Infected: 5
['0' 'Z' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' 'Z' '0' '0'
 '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0'
 '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' '0' 'Z' 'Z'
 '0' '0' '0' 'Z' '0' '0']
Number Infected: 0

问题原因 #

1. 随机数生成长度错误：numpy.random.binomial(1,p*count)未指定size参数，默认只生成1个随机值，所有原0的位置都被替换成同一个值（本次运行刚好是0），自然看不到新的1。
2. 类型不匹配导致统计失败：替换"Z"后数组 dtype 变为字符串类型，新生成的1会被转成"1"，后续判断ele == 1（整数）无法匹配，所以count为0。

修复后的代码 #

import numpy

n = 50  # 运行重复次数
p = 0.1

# 模拟第一天
result = numpy.random.binomial(1, p, 60)
print(result)
count = numpy.sum(result == 1)  # 用numpy内置求和替代循环，效率更高
print("Number Infected:", count)

# 创建追踪感染情况的矩阵
# 生成与原数组长度一致的随机数组
new_random = numpy.random.binomial(1, p*count, size=60)
# 执行替换
result = numpy.where(result == 1, "Z", new_random)
print(result)

# 统计感染数：同时匹配整数1和字符串"1"
count = numpy.sum((result == 1) | (result == "1"))
print("Number Infected:", count)

额外优化建议 #

如果不想混合字符串和整数类型，可以用整数标记状态（比如0=未感染，1=新感染，2=已标记为Z），最后再统一转换为字符串，这样类型更统一，避免判断逻辑混乱。

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