python - 如何计算 Pandas 数据框中到周末或休息日的天数

Question

我有一个带有非连续日期索引的 pandas 数据框(缺少周末和节假日)。我想添加包含到下一天休息的天数的列。

这里是生成示例数据框的代码，在 till_day_off 列中包含所需的值:

import pandas as pd
​
df = pd.DataFrame(index=pd.date_range(start="2022-06-06", periods=15))
df["day_of_week"] = df.index.dayofweek   # adding column with number of day in a week
df = df[(df.day_of_week < 5)]   # remove weekends
df = df.drop(index="2022-06-15")   # remove Wednesday in second week
df["till_day_off"] = [5,4,3,2,1,2,1,2,1,1] # desired values, end of column is treated as day off

结果数据框:

<表类="s-表"><头><日> day_of_weektill_day_off<正文>2022-06-06052022-06-07142022-06-08232022-06-09322022-06-10412022-06-13022022-06-14112022-06-16322022-06-17412022-06-2001

真实数据框有超过 7_000 行，所以显然我试图避免对行进行迭代。知道如何解决这个问题吗？

Answer 1

假设输入已排序(如果没有，则按天排序)，您可以使用掩码来识别连续的天数并使用它对它们进行分组并计算累计数:

mask = (-df.index.to_series().diff(-1)).eq('1d').iloc[::-1]
# reversing the Series to count until (not since) the value

df['till_day_off'] = mask.groupby((~mask).cumsum()).cumcount().add(1)

输出:

            day_of_week  till_day_off
2022-06-06            0             5
2022-06-07            1             4
2022-06-08            2             3
2022-06-09            3             2
2022-06-10            4             1
2022-06-13            0             2
2022-06-14            1             1
2022-06-16            3             2
2022-06-17            4             1
2022-06-20            0             1

中间体:

mask

2022-06-20    False
2022-06-17    False
2022-06-16     True
2022-06-14    False
2022-06-13     True
2022-06-10    False
2022-06-09     True
2022-06-08     True
2022-06-07     True
2022-06-06     True
dtype: bool

(~mask).cumsum()

2022-06-20    1
2022-06-17    2
2022-06-16    2
2022-06-14    3
2022-06-13    3
2022-06-10    4
2022-06-09    4
2022-06-08    4
2022-06-07    4
2022-06-06    4
dtype: int64