如何基于日期差条件填充Pandas数据框中的NA值？

问题 #

需要在Python中实现以下条件式空值替换操作：

• 当month列与Date1列的月份差≤2时，将num列中的空值（np.nan）替换为0；差值大于2则不替换。
• 当month列与Date2列的月份差≤2时，将Num1、Num2列中的空值替换为0；差值大于2则不替换。

生成数据集的代码：

import pandas as pd
import numpy as np

data = {'month': ['2022-01-01', '2022-02-01', '2022-03-01', '2022-01-01', '2022-02-01', '2022-03-01', '2022-04-01', '2022-05-01', '2022-06-01', '2022-07-01', '2022-08-01'], 'Date1': ['2022-01-01', '2022-01-01', '2022-01-01', '2022-01-01', '2022-01-01', '2022-01-01', '2022-01-01', '2022-05-01', '2022-05-01', '2022-05-01', '2022-05-01'], 'Date2': ['2022-02-01', '2022-02-01', '2022-02-01', '2022-04-01', '2022-04-01', '2022-04-01', '2022-04-01', np.nan, np.nan, np.nan, np.nan], 'Name': ['A', 'A', 'A', 'B', 'B', 'B', 'B', 'C', 'C', 'C', 'C'], 'num': [1234, 1234, 1234, 456, 456, 456, 456, np.nan, np.nan, np.nan, np.nan], 'sales': ['MN', 'MN', 'MN', 'CA', 'CA', 'CA', 'CA', 'FL', 'FL', 'FL', 'FL'], 'Num1': [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 44.0, 44.0, 44.0, 44.0], 'Num2': [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, 29.0, 29.0, 29.0, 29.0]}

df = pd.DataFrame(data)

初始数据框内容：

month       Date1       Date2 Name   num sales  Num1  Num2
0   2022-01-01  2022-01-01  2022-02-01    A  1234    MN   NaN   NaN
1   2022-02-01  2022-01-01  2022-02-01    A  1234    MN   NaN   NaN
2   2022-03-01  2022-01-01  2022-02-01    A  1234    MN   NaN   NaN
3   2022-01-01  2022-01-01  2022-04-01    B   456    CA   NaN   NaN
4   2022-02-01  2022-01-01  2022-04-01    B   456    CA   NaN   NaN
5   2022-03-01  2022-01-01  2022-04-01    B   456    CA   NaN   NaN
6   2022-04-01  2022-01-01  2022-04-01    B   456    CA   NaN   NaN
7   2022-05-01  2022-05-01         NaN    C   NaN    FL  44.0  29.0
8   2022-06-01  2022-05-01         NaN    C   NaN    FL  44.0  29.0
9   2022-07-01  2022-05-01         NaN    C   NaN    FL  44.0  29.0
10  2022-08-01  2022-05-01         NaN    C   NaN    FL  44.0  29.0

解决方案 #

步骤1：转换日期格式并计算月份差 #

先将日期列转为datetime类型，再计算month与Date1、Date2的绝对月份差：

# 转换日期格式
df['month'] = pd.to_datetime(df['month'])
df['Date1'] = pd.to_datetime(df['Date1'])
df['Date2'] = pd.to_datetime(df['Date2'])

# 计算月份差：年份差*12 + 月份差，取绝对值
df['month_diff1'] = abs((df['month'].dt.year - df['Date1'].dt.year)*12 + (df['month'].dt.month - df['Date1'].dt.month))
df['month_diff2'] = abs((df['month'].dt.year - df['Date2'].dt.year)*12 + (df['month'].dt.month - df['Date2'].dt.month))

步骤2：替换num列空值 #

根据month_diff1的条件，替换符合要求的空值：

df['num'] = np.where((df['month_diff1'] <= 2) & (df['num'].isna()), 0, df['num'])

步骤3：替换Num1、Num2列空值 #

根据month_diff2的条件，替换这两列的空值：

df['Num1'] = np.where((df['month_diff2'] <= 2) & (df['Num1'].isna()), 0, df['Num1'])
df['Num2'] = np.where((df['month_diff2'] <= 2) & (df['Num2'].isna()), 0, df['Num2'])

步骤4：（可选）移除临时列 #

如果不需要保留中间计算的月份差列，可删除：

df = df.drop(['month_diff1', 'month_diff2'], axis=1)

最终结果 #

处理后的数据框如下：

month      Date1      Date2 Name   num sales  Num1  Num2
0   2022-01-01 2022-01-01 2022-02-01    A  1234    MN   0.0   0.0
1   2022-02-01 2022-01-01 2022-02-01    A  1234    MN   0.0   0.0
2   2022-03-01 2022-01-01 2022-02-01    A  1234    MN   0.0   0.0
3   2022-01-01 2022-01-01 2022-04-01    B   456    CA   NaN   NaN
4   2022-02-01 2022-01-01 2022-04-01    B   456    CA   NaN   NaN
5   2022-03-01 2022-01-01 2022-04-01    B   456    CA   0.0   0.0
6   2022-04-01 2022-01-01 2022-04-01    B   456    CA   0.0   0.0
7   2022-05-01 2022-05-01        NaT    C     0    FL  44.0  29.0
8   2022-06-01 2022-05-01        NaT    C     0    FL  44.0  29.0
9   2022-07-01 2022-05-01        NaT    C   NaN    FL  44.0  29.0
10  2022-08-01 2022-05-01        NaT    C   NaN    FL  44.0  29.0

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