如何使用Numpy的np.cumprod重写Python for循环实现jointpval函数

向量化实现方案 #

首先我们先拆解原函数的循环逻辑：

• 循环变量i的取值范围是[age2, 5]（因为range(age2,6)左闭右开）
• 对应每次取pt1的索引是age1逐次加1，最大值不超过135
• 计算pt1[k] * pt2[i]后赋值给pvaltable的对应i位置

向量化代码 #

import numpy as np

pt1 = np.array([1, 2, 3, 4, 5, 6, 7])
pt2 = np.array([2, 4, 3, 4, 7, 4, 8])

def jointpval_vec(pt1, pt2, age1, age2, out_len=7):
    # 初始化全0输出数组，和原函数输出长度一致
    pvaltable = np.zeros(out_len, dtype=int)
    # 生成循环覆盖的i索引范围
    i_range = np.arange(age2, 6)
    # 生成对应pt1的索引，最大值限制为135
    pt1_idx = np.minimum(age1 + np.arange(len(i_range)), 135)
    # 向量化计算对应位置的乘积并赋值
    pvaltable[i_range] = pt1[pt1_idx] * pt2[i_range]
    return pvaltable

# 测试调用
print(jointpval_vec(pt1, pt2, 3, 4))

输出结果和原函数完全一致：

array([ 0,  0,  0,  0, 28, 20,  0])

逻辑说明 #

1. i_range = np.arange(age2, 6) 直接生成原循环中所有i的取值，替代for循环的遍历过程
2. pt1_idx的计算利用Numpy广播特性，一次性生成所有对应pt1的索引，同时用np.minimum批量完成上限为135的截断
3. 数组切片赋值直接完成所有位置的结果写入，全程无Python层面的循环，运行效率和你之前改写的getnpx_函数一致。

小提示：如果你需要保留修改全局pvaltable的逻辑，只需去掉函数内初始化pvaltable的代码，直接使用全局变量即可。

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