numpy向量非零元素求逆问题：vectorize异常及高效解法问询

np.vectorize 并非真正的向量化实现，本质是对数组元素逐个调用函数的循环包装，且会根据第一个元素的返回值类型推断整个输出数组的数据类型（dtype）。

你的代码中，第一个元素是0，lambda返回整数0，vectorize因此将输出数组的dtype设为整数类型。后续非零元素计算得到的浮点数（比如1./1=1.0）会被强制转换为整数，结果变为0，最终整个数组全是0。

方法1：指定vectorize的输出类型 #

显式设置otypes参数为浮点数类型，避免类型推断错误：

import numpy as np
active_N = np.array([0,1,3,5])
f = np.vectorize(lambda x: x if x==0 else 1./x, otypes=[np.float64])
active_N_inverse = f(active_N)
# 输出：array([0.        , 1.        , 0.33333333, 0.2       ])

方法2：使用numpy原生向量化操作（更高效） #

np.vectorize性能和普通循环相近，推荐用numpy原生向量化逻辑直接操作数组，速度更快：

import numpy as np
active_N = np.array([0,1,3,5])
active_N_inverse = np.zeros_like(active_N, dtype=np.float64)
# 对非零元素求逆
mask = active_N != 0
active_N_inverse[mask] = 1. / active_N[mask]
# 输出：array([0.        , 1.        , 0.33333333, 0.2       ])

也可以用更简洁的np.where写法：

active_N_inverse = np.where(active_N != 0, 1. / active_N, 0.)

这两种原生方法是真正的向量化运算，处理大规模数组时，效率远高于np.vectorize。

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