使用np.vectorize装饰含while循环的函数时循环无法正常终止

问题复现 #

你的代码在移除@np.vectorize后可正常运行，但添加装饰器后，while循环触发break后似乎仍继续执行，需两次break才停止：

import numpy as np
from numpy import random

@np.vectorize
def simulation(_):
    numRolls = 0
    diff = 0
    prevRoll = -999999
    while True:
        roll = random.randint(1, 7)
        print(prevRoll, roll)
        numRolls += 1
        diff = np.abs(roll - prevRoll)
        if diff == 1:
            break
        else:
            prevRoll = roll
    
    return numRolls

原因分析 #

np.vectorize并非真正的向量化运算，它本质是一个循环包装器——帮你把处理标量的函数批量应用到数组元素上。关键问题在于：
当你未指定otypes参数时，np.vectorize会自动调用一次你的函数（通常用输入的第一个元素或默认值），目的是推断输出结果的数据类型。这次试探性调用会完整执行你的while循环直到break，之后才会执行你实际需要的调用。
你看到的“两次break”，其实是两次独立的函数调用（一次试探、一次实际执行），而非同一个循环未停止。

解决方案 #

方案1：指定otypes参数跳过试探调用 #

直接在装饰器中声明函数返回值的类型，避免np.vectorize额外调用函数推断类型：

import numpy as np
from numpy import random

@np.vectorize(otypes=[np.int32])  # 指定返回值为int32类型
def simulation(_):
    numRolls = 0
    diff = 0
    prevRoll = -999999
    while True:
        roll = random.randint(1, 7)
        numRolls += 1
        diff = np.abs(roll - prevRoll)
        if diff == 1:
            break
        prevRoll = roll
    
    return numRolls

# 调用示例：模拟1000次并计算均值
n_simulations = 1000
results = simulation(np.zeros(n_simulations))  # 用全0数组触发n_simulations次调用
mean_rolls = np.mean(results)
print(mean_rolls)

方案2：放弃np.vectorize，直接用列表推导式+np.mean #

既然np.vectorize本质还是循环，不如直接用更直观的列表推导式生成模拟结果，再计算均值，效率几乎一致：

import numpy as np
from numpy import random

def simulation(_):
    numRolls = 0
    diff = 0
    prevRoll = -999999
    while True:
        roll = random.randint(1, 7)
        numRolls += 1
        diff = np.abs(roll - prevRoll)
        if diff == 1:
            break
        prevRoll = roll
    
    return numRolls

# 模拟1000次并计算均值
n_simulations = 1000
results = [simulation(_) for _ in range(n_simulations)]
mean_rolls = np.mean(results)
print(mean_rolls)

方案3：真正的向量化模拟（高效版） #

如果需要处理大规模模拟，上述循环方式效率较低，可以用numpy的数组操作实现完全向量化的模拟，避免Python层面的循环：

import numpy as np

def vectorized_simulation(n_simulations):
    # 预生成足够多的随机骰子数（每个模拟最多预设20次roll，不够再补）
    max_rolls = 20
    rolls = np.random.randint(1, 7, size=(n_simulations, max_rolls))
    # 计算相邻roll的差的绝对值
    diffs = np.abs(rolls[:, 1:] - rolls[:, :-1])
    # 找到每个模拟中第一个diff==1的位置
    first_break = np.argmax(diffs == 1, axis=1)
    # 处理从未触发break的情况（概率极低，这里直接补max_rolls）
    first_break[diffs.sum(axis=1) == 0] = max_rolls - 1
    # 总roll次数是第一个break位置+2（因为从第2个roll开始算diff，加上初始的1次roll）
    total_rolls = first_break + 2
    return total_rolls

# 模拟10000次并计算均值
n_simulations = 10000
results = vectorized_simulation(n_simulations)
mean_rolls = np.mean(results)
print(mean_rolls)

内容的提问来源于stack exchange，提问作者MAN-MADE