python - 如何将 pandas 中数字后面的 N 个索引设为空？

Question

我有一个包含值和 NaN 的 Series，如下所示:

0    NaN
1    1.0
2    NaN
3    2.0
4    3.0
5    NaN
6    NaN
7    4.0
8    5.0
9    NaN
dtype: float64

举例来说，我想在每个初始值之后null出接下来的四个索引，如下所示:

0    NaN
1    1.0
2    NaN
3    NaN
4    NaN
5    NaN
6    NaN
7    4.0
8    NaN
9    NaN
dtype: float64

(第一个 1.0 之后的四个值为空，以及下一个值 4.0 之后的最后两个为空。)

Answer 1

方法#1

这是一种循环方式，仅迭代非空位置列表 -

def nullnext(s, W):
    a = s.values
    idx = np.flatnonzero(s.notnull().values)+1
    last_idx = idx[0]
    a[last_idx:last_idx+W] = np.nan
    for i in idx[1:]:
        if i > last_idx + W:
            last_idx = i
            a[last_idx:last_idx+W] = np.nan
    return s

示例运行 -

In [336]: s
Out[336]: 
0    1.0
1    NaN
2    2.0
3    3.0
4    NaN
5    NaN
6    4.0
7    5.0
8    NaN
Name: NaN, dtype: float64

In [337]: nullnext(s, W=4)
Out[337]: 
0    1.0
1    NaN
2    NaN
3    NaN
4    NaN
5    NaN
6    4.0
7    NaN
8    NaN
Name: NaN, dtype: float64

方法#2

只需进行一些调整，我们就可以将其移植到 numba 上以提高性能效率。实现涉及使用strides。相关代码看起来像这样 -

from numba import njit

# https://stackoverflow.com/a/40085052/ @Divakar
def strided_app(a, L, S ):  # Window len = L, Stride len/stepsize = S
    nrows = ((a.size-L)//S)+1
    n = a.strides[0]
    return np.lib.stride_tricks.as_strided(a, shape=(nrows,L), strides=(S*n,n))

@njit
def set_mask(mask, idx, W):
    last_idx = idx[0]
    mask[0] = True
    l = len(idx)
    for i in range(1,l):
        if idx[i] > last_idx + W:
            last_idx = idx[i]
            mask[i] = True
    return mask


def nullnext_numba(s, W):
    a = s.values
    idx = np.flatnonzero(s.notnull().values)+1

    mask = np.zeros(len(idx),dtype=bool)
    set_mask(mask, idx, W)

    a_ext = np.concatenate((a, np.full(W,np.nan)))
    strided_app(a_ext, W, 1)[idx[mask]] = np.nan
    return pd.Series(a_ext[:-W])

进一步改进

我们可以进一步优化它，通过避免串联来提高内存效率，并使用输入序列进行所有这些编辑，从而提高性能，就像这样 -

def nullnext_numba_v2(s, W):
    a = s.values
    idx = np.flatnonzero(s.notnull().values)+1

    mask = np.zeros(len(idx),dtype=bool)
    set_mask(mask, idx, W)

    valid_idx = idx[mask]    
    limit_mask = valid_idx < len(a) - W
    strided_app(a, W, 1)[valid_idx[limit_mask]] = np.nan

    leftover_idx = valid_idx[~limit_mask]
    if len(leftover_idx)>0:
        a[leftover_idx[0]:] = np.nan
    return s