python - 仅当间隙小于 N nan 时才填充 pandas df 中的 nan 间隙

Question

我正在使用还包含 nan 值的 pandas 数据框。我想用 df.interpolate 替换带有插值的 nan，但前提是 nan 值序列的长度为 =

print(df)
A   B   C
1   1   1
nan nan 2
nan nan 3
nan 4   nan
5   5   5

在这种情况下，我想在 df 上应用一个函数，仅填充长度为 N<=2 的 nan 序列，但较大的序列保持不变，从而得到我想要的输出

print(df)
A   B   C
1   1   1
nan 2   2
nan 3   3
nan 4   4
5   5   5

请注意，我知道 df.interpolate 中的 limit=N 选项，但它不能满足我的要求，因为它会填充任何长度对于 nan 序列，只需将填充限制为前 3 个 nan，从而导致不需要的输出

print(df)
A   B   C
1   1   1
2   2   2
3   3   3
nan 4   4
5   5   5

那么你知道一个函数吗/你知道如何构造一个代码来产生我想要的输出吗？ TNx

Answer 1

您可以执行游程长度编码并识别每列小于或等于两个元素的 NaN 游程。一种方法是使用 pdrle 包中的 get_id (免责声明:我写的)。

import pdrle


chk = df.isna() & (df.apply(lambda x: x.groupby(pdrle.get_id(x)).transform(len)) <= 2)
df[chk] = df.interpolate()[chk]
#      A    B    C
# 0  1.0  1.0  1.0
# 1  NaN  2.0  2.0
# 2  NaN  3.0  3.0
# 3  NaN  4.0  4.0
# 4  5.0  5.0  5.0