python - numpy 前向填充条件

Question

我有一个像这样的带有零的 numpy 数组。

a = np.array([3., 0., 2., 3., 0., 3., 3., 3., 0., 3., 3., 0., 3., 0., 0., 0., 0.,
   3., 3., 0., 3., 3., 0., 3., 0., 3., 0., 0., 0., 3., 0., 3., 3., 0.,
   3., 3., 0., 0., 3., 0., 0., 0., 3., 0., 3., 3., 3., 3., 3., 3., 3.,
   3., 3., 3., 3., 3., 3., 4., 3., 0., 3., 3., 3., 3., 3., 3., 3., 0.,
   0., 0., 0., 3., 0., 0., 3., 0., 0., 0., 3., 3., 3., 3., 3., 3., 3.,
   3., 0., 3., 3., 3., 3., 3., 0., 3., 3., 3., 3., 0., 0., 0., 3., 3.,
   3., 0., 3., 3., 3., 5., 3., 3., 3., 3., 3., 3., 3., 0., 3., 0., 3.,
   3., 0., 0., 0., 3., 3., 3., 3., 0., 3., 3., 3., 3., 3., 3., 3., 3.,
   3., 3., 3., 3., 0., 3., 3., 3., 3., 3., 3., 0., 3., 3., 3., 3., 3.,
   3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 0., 3., 0., 3.,
   3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 3., 0., 3., 3., 3., 3.,
   3., 3., 3., 3., 3., 3., 3., 3., 0., 3., 3., 0., 0., 3., 0., 0., 3.,
   0., 3., 3., 0., 3., 3., 0., 0., 3., 3., 3., 3., 3., 3., 3., 0., 3.,
   3., 3., 3., 3.])

我需要在某个条件下用以前的值(前向填充)替换零。如果两个非零数字之间的零数小于或等于 2，则需要前向填充零。

举个例子，

1) 如果我考虑3., 0., 2.这三个数字，非零数字之间的零个数为1。这应该用3填充。

2) 如果我考虑 3., 0., 0., 0., 0.,3., 3. 这些数字，3 之间的零数大于 2.so将保持原样。

Answer 1

在这些情况下，提出纯矢量化方法似乎并不简单(至少在这种情况下可以这么说)，我们可以使用 numba 将您的代码编译为 C级别。这是使用 numba 的 nopython 模式的一种方法:

import numba

@numba.njit('int64[:](int64[:],uintc)') #change accordingly
def conditional_ffill(a, w):
    c=0
    last_non_zero = a[0]
    out = np.copy(a)
    for i in range(len(a)):
        if a[i]==0:
            c+=1
        elif c>0 and c<w:
            out[i-c:i] = last_non_zero
            c=0
            last_non_zero=a[i]
    return out

---

检查 divakar 的测试数组:

a = np.array([2, 0, 3, 0, 0, 4, 0, 0, 0, 5, 0])

conditional_ffill(a, w=1)
# array([2, 0, 3, 0, 0, 4, 0, 0, 0, 5, 0])

conditional_ffill(a, w=2)
# array([2, 2, 3, 0, 0, 4, 0, 0, 0, 5, 0])

conditional_ffill(a, w=3)
# array([2, 2, 3, 3, 3, 4, 0, 0, 0, 5, 0])

conditional_ffill(a, w=4)
# array([2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 0])

更大阵列上的计时:

a_large = np.tile(a, 10000)

%timeit ffill_windowed(a_large, 3)
# 1.39 ms ± 68.7 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit conditional_ffill(a_large, 3)
# 150 µs ± 862 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)