如何调整Newton-Raphson的for循环，传入多组初值并输出每组迭代次数

• 原代码使用math.exp仅支持单标量计算，传入numpy数组时会触发类型不匹配错误
• 原迭代逻辑的收敛判断为全局单条件break，多初值场景下不同样本收敛速度不同，无法直接沿用单标量的跳出规则

只需要替换数学库为numpy的向量化实现、调整收敛判断逻辑即可，不需要重构核心迭代结构：

import numpy as np

# 适配数组输入的Newton-Raphson方法
def nraphson(fn, dfn, x, tol, maxiter):
    # 初始化迭代次数数组、收敛掩码
    iter_counts = np.zeros_like(x, dtype=int)
    converged = np.zeros_like(x, dtype=bool)
    x_current = x.copy()
    
    for i in range(maxiter):
        # 仅更新未收敛的样本
        x_new = np.empty_like(x_current)
        x_new[converged] = x_current[converged]
        x_new[~converged] = x_current[~converged] - fn(x_current[~converged])/dfn(x_current[~converged])
        
        # 判断新收敛的样本，记录迭代次数
        newly_converged = abs(x_new - x_current) < tol
        iter_counts[newly_converged & (~converged)] = i + 1
        converged[newly_converged] = True
        
        x_current = x_new
        # 所有样本都收敛则提前退出
        if np.all(converged):
            break
    # 未收敛的样本迭代次数记为最大迭代次数
    iter_counts[~converged] = maxiter
    return x_current, iter_counts

# 替换math.exp为np.exp，支持数组计算
y = lambda x: np.exp(x) - x**2
dy = lambda x: np.exp(x) - 2*x

# 测试多初值输入，比如-5到0步长0.5的所有初值
init_x = np.arange(-5, 0, 0.5)
roots, iters = nraphson(y, dy, init_x, 0.0001, 100)

# 打印每个初值对应的结果
for init, root, n in zip(init_x, roots, iters):
    print(f"初值{init:.1f}，根为{root:.6f}，迭代次数{n}")

• 输入初值可以是任意形状的numpy数组，返回的根数组、迭代次数数组和输入形状完全一致
• 已收敛的初值不会再参与后续迭代计算，性能和单初值迭代基本持平
• 未在最大迭代次数内收敛的样本，迭代次数会统一标记为设置的maxiter值，方便后续筛选异常初值

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