python - 是否有一种快速的方法来填充间隙(可能是多个)之间的 NA 值，并且仅当间隙小于一定大小时？

Question

我有一个带有 id 和 date 的 DataFrame。我想创建另一个 DataFrame ，它列出 id 是否存在于给定的月份，超过预定义的时间段(比如 2018 年全年)。此外，如果有任何 2 个月或更短的存在差距，我想填补它们。

我想我会包括第一部分，因为从一开始就有更好的解决方案。这是开始的 df

import pandas as pd
import numpy as np
df = pd.DataFrame({'id': [1,1,1,1,1,2,2,2,3],
                 'date': ['2018-02-01', '2018-03-12', '2018-05-10', 
                         '2018-10-10', '2018-11-04', '2018-06-07', '2018-07-07', 
                         '2018-09-16', '2018-02-02']})
df['date'] = pd.to_datetime(df.date)

为了让它存在 df_exist 我创建了一个新列 id_exists 并与平铺时间段合并 df_per

df['id_exists'] = True
per = pd.date_range('2018-01-01', '2018-12-31', freq='MS')

df_per = pd.DataFrame({'id': np.tile(df.id.unique(), len(per)),
    'Period': np.repeat(per,df.id.nunique())})
df_exist = df_per.merge(df, left_on=['id', df_per.Period.dt.year, df_per.Period.dt.month],
             right_on=['id', df.date.dt.year, df.date.dt.month], how='left').drop(columns='date').fillna(False)

#      Period  id  id_exists
#0 2018-01-01   1      False
#1 2018-01-01   2      False
#2 2018-01-01   3      False
#3 2018-02-01   1       True
#4 2018-02-01   2      False

我决定用 False 来 fillna 因为这允许我使用下面的函数和 cumsum，但当然如果有使用 NaN 的解决方案也一样好。

现在我已经定义了一个函数，它似乎可以做我想做的事情:索引比较确保我不会在任何一边填充东西，而与 gap_size 的比较确保我只填充小的空白。 if-else 确保它正常工作，无论第一个条目在存在 df 中是 True 还是 False。

def FillGaps(df, gap_size):
    gb = df.groupby(df.id_exists.cumsum()).size()

    if df.id_exists.values[0] == False:
        to_fill = gb[(gb.index > gb.index.min()) & (gb.index < gb.index.max()) &
            (gb.values <= gap_size)].index.values
    else:
        to_fill = gb[(gb.index < gb.index.max()) & (gb.values <= gap_size)].index.values

    df.loc[df.id_exists.cumsum().isin(to_fill), 'id_exists'] = True
    return df

df_exist = df_exist.groupby('id').apply(lambda df: FillGaps(df, gap_size=2))

但是，它在大型 DataFrame 上非常慢。关于如何使它更快的任何想法？它看起来不像任何内置的 fillna 方法都适用于这种类型的填补空白的情况，可能存在多个空白。

这是预期的输出。 (我做了一些合并，所以它没有格式化为一个烦人的长表)。关键是没有任何边缘受到干扰，只有 2 个月或更短的间隙被填充并且 id==3 没有失败，其中只有一个值开始。

       Period  id_1  id_exists_1  id_2  id_exists_2  id  id_exists
0  2018-01-01     1        False     2        False   3      False
1  2018-02-01     1         True     2        False   3       True
2  2018-03-01     1         True     2        False   3      False
3  2018-04-01     1         True     2        False   3      False
4  2018-05-01     1         True     2        False   3      False
5  2018-06-01     1        False     2         True   3      False
6  2018-07-01     1        False     2         True   3      False
7  2018-08-01     1        False     2         True   3      False
8  2018-09-01     1        False     2         True   3      False
9  2018-10-01     1         True     2        False   3      False
10 2018-11-01     1         True     2        False   3      False
11 2018-12-01     1        False     2        False   3      False

Answer 1

这是一种实现方式:

month = df.date - pd.Timedelta('1 day') * (df.date.dt.day - 1)
df_exist = df.id.astype(str).str.get_dummies().groupby(month).sum() != 0

def fill_gaps(arr):
    notnan, = (~np.isnan(arr)).nonzero()
    return np.nan if not notnan.size else arr[notnan[-1]]

date_range = pd.date_range('2018-01-01', '2018-12-31', freq='MS')
rolling = df_exist.reindex(date_range).rolling(window=2, min_periods=1)
result = rolling.apply(fill_gaps).fillna(False).astype(bool)
result[date_range > month.max()] = False

第一部分应该比手动连接快得多。第二部分对数据帧使用滚动 API。

输出看起来像这样:

                1      2      3
2018-01-01  False  False  False
2018-02-01   True  False   True
2018-03-01   True  False  False
2018-04-01   True  False  False
2018-05-01   True  False  False
2018-06-01  False   True  False
2018-07-01  False   True  False
2018-08-01  False   True  False
2018-09-01  False   True  False
2018-10-01   True  False  False
2018-11-01   True  False  False
2018-12-01  False  False  False

这似乎与您示例中的结果相匹配。