Pandas如何基于非周末DatetimeIndex获取指定日期的下一有效日期

问题描述 #

我有两个变量 a 和 b：a 为 DataFrame 类型，b 为 Series 类型。a.index 包含2020年绝大多数非周末日期，b.index 为我随机选取的若干日期，相关代码及输出如下：

import pandas as pd
from datetime import datetime
print(a.index)
DatetimeIndex(['2020-01-02', '2020-01-03', '2020-01-06', '2020-01-07',
               '2020-01-08', '2020-01-09', '2020-01-10', '2020-01-13',
               '2020-01-14', '2020-01-15',
               ...
               '2020-12-18', '2020-12-21', '2020-12-22', '2020-12-23',
               '2020-12-24', '2020-12-25', '2020-12-28', '2020-12-29',
               '2020-12-30', '2020-12-31'],
              dtype='datetime64[ns]', name='Date', length=245, freq=None)


print(b.index)
DatetimeIndex(['2020-02-03', '2020-03-02', '2020-03-04', '2020-03-10',
               '2020-03-13'],
              dtype='datetime64[ns]', name='Date', freq=None)

我想要获取 b.index 中每个日期在 a.index 内对应的下一个存在的日期，最初的尝试代码如下：

target_dates = b.index + datetime.timedelta(days=1)
a.loc[target_dates]

但运行后抛出 KeyError: "[Timestamp('2020-03-14 00:00:00')] not in index" 错误，原因是2020-03-14为周末，不在 a.index 范围内。我期望得到的最终结果如下：

DatetimeIndex(['2020-02-04', '2020-03-03', '2020-03-05', '2020-03-11',
               '2020-03-16'],
              dtype='datetime64[ns]', name='Date', freq=None)

请问如何通过 a.index 和 b.index 生成上述符合要求的 DatetimeIndex？

解决方案 #

以下是两种可直接实现需求的方案：

方案1：使用 searchsorted（性能最优，适合大数据量）

searchsorted 可以快速定位目标值在有序序列中的插入位置，正好对应第一个大于等于目标值的元素位置，使用前需确保 a.index 为升序排列：

# 若a的索引已经是升序，可跳过排序步骤
a = a.sort_index()
# 计算每个b的日期加1天后，在a.index中对应的插入位置
positions = a.index.searchsorted(b.index + pd.Timedelta(days=1))
# 按位置提取对应的日期
target_dates = a.index[positions]

方案2：使用 reindex 后向填充（写法更直观）

利用 reindex 的 bfill 填充规则，自动匹配目标日期之后最近的有效索引值：

target_dates = a.index.reindex(b.index + pd.Timedelta(days=1), method='bfill')[0]

两种方案输出的 target_dates 都和预期结果完全一致。

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