为何np.zeros_like(x)与np.zeros(len(x))生成的数组结果不同？

问题复现 #

使用np.zeros(len(x))时函数能得到正确浮点结果：

import numpy as np

def discount_cumsum(x: np.ndarray, d):
    # 该行代码可正常运行
    y = np.zeros(len(x))

    y[-1] = x[-1]
    for i in range(len(x) - 2, -1, -1):
        y[i] = x[i] + d * y[i + 1]
    return y

discount_cumsum(np.array(range(5)), 0.5) # array([1.625, 3.25 , 4.5  , 5.   , 4.   ])

切换为np.zeros_like(x)时得到错误的整数结果：

def discount_cumsum_2(x: np.ndarray, d):
    # 该行代码无法得到正确结果
    y = np.zeros_like(x)

    y[-1] = x[-1]
    for i in range(len(x) - 2, -1, -1):
        y[i] = x[i] + d * y[i + 1]
    return y

discount_cumsum_2(np.array(range(5)), 0.5) # array([1, 3, 4, 5, 4])

差异核心原因：数据类型不一致 #

两者的本质区别在于生成数组的数据类型：

• np.zeros(len(x))默认生成float64类型的浮点数组，计算过程中带小数的中间结果能完整保留，最终输出正确的浮点结果。
• np.zeros_like(x)会完全继承输入x的数据类型。这里传入的range(5)转为numpy数组后是int64整数类型，因此生成的y也是整数数组。计算时，所有带小数的结果会被自动截断取整（比如2.5直接转为2），最终得到错误的整数输出。

如果要让np.zeros_like生成浮点数组，只需显式指定数据类型：

y = np.zeros_like(x, dtype=np.float64)

正确实现参考 #

你最初的代码通过显式转换数据类型避免了这个问题，是更严谨的写法：

import numpy as np
from numba import njit

@njit
def discount_cumsum(x: np.ndarray, d):
    y = x.copy().astype(np.float32)
    for i in range(len(y) - 2, -1, -1):
        y[i] += d * y[i + 1]
    return y

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