如何用纯numpy实现数组nan值前向填充最多N次超出补0

需求规则 #

• 使用最近的非nan值填充后续连续nan，最多填充N次
• 若连续nan数量大于N，超出部分统一用0填充
• 不得使用迭代循环，仅调用纯numpy接口实现

示例参考 #

测试用例代码：

import numpy as np

n = 2
arr = np.array([np.nan, 0, 0, np.nan, 5, 4, 4, np.nan, np.nan, np.nan, 1, 5, 3, np.nan, 2, np.nan, np.nan])

def ffill(arr: np.array, n: int):
    pass
    return arr

预期输出结果：

result = np.array([0.0, 0.0, 0.0, 0.0, 5.0, 4.0, 4.0, 4.0, 4.0, 0.0, 1.0, 5.0, 3.0, 3.0, 2.0, 2.0, 2.0])

示例说明：数值4后续的3个连续nan在n=2的限制下，前2个填充为4，最后1个超出限制填充为0，即[... 4, np.nan, np.nan, np.nan ...] 填充后为 [... 4, 4, 4, 0 ...]；数组开头的nan默认前置填充为0。

实现方案（适配大N场景优化） #

def ffill(arr: np.array, n: int):
    # 开头无前置有效值的nan默认填充0
    if np.isnan(arr[0]):
        arr[0] = 0

    isnan = np.isnan(arr)
    notnan = ~isnan
    # 提取所有非nan有效值
    valid = arr[notnan]
    # 生成前向填充的映射索引
    indices = np.cumsum(notnan) - 1
    arr = valid[indices]

    # 标记超出最大填充次数的nan位置
    overlimit = np.lib.stride_tricks.sliding_window_view(isnan[:-1][::-1], isnan.size-n)[:, ::-1].all(axis=0)
    overlimit &= isnan[n:]

    # 超出限制的位置统一置0
    indices = np.flatnonzero(overlimit) + n
    arr[indices] = 0
    return arr

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