基于日期范围匹配为Pandas DataFrame生成id与parent字段值

需求说明

• 若某条记录的日期范围完全处于某已赋值父记录的日期范围内，则其id为父记录id追加序号（如A-1.1），parent为父记录id；
• 若未匹配到符合条件的父记录，则分配新的顶级id（如A-2），parent设为A-Root。

输入DataFrame

>>> mydf
            name  start_date    end_date      id  parent
0       Total data 2020-01-01 2021-03-18  A-Root    None
1          Husband 2020-01-01 2021-03-15     A-1  A-Root
2    Not-in-family 2020-01-02 2021-03-18       -       -
3        Unmarried 2020-01-04 2021-03-04       -       -
4             Wife 2020-01-10 2021-03-17       -       -
5        Own-child 2020-01-11 2021-03-16       -       -

预期输出DataFrame

name start_date   end_date      id  parent
0       Total data 2020-01-01 2021-03-18  A-Root    None
1          Husband 2020-01-01 2021-03-15     A-1  A-Root
2    Not-in-family 2020-01-02 2021-03-18     A-2  A-Root
3        Own-child 2020-01-11 2021-03-16   A-2.1     A-2
4        Unmarried 2020-01-04 2021-03-04   A-1.1     A-1
5             Wife 2020-01-10 2021-03-17     A-3     A-2

解决方案

思路 #

1. 先将已完成赋值的记录存入字典，记录每个id对应的日期范围；
2. 遍历未赋值的记录，逐个检查是否能匹配已有的父记录日期范围；
3. 匹配成功则生成带序号的子id，未匹配则生成新的顶级id，同时更新字典记录新的赋值信息。

代码实现 #

import pandas as pd

# 转换日期列为datetime格式，方便比较
mydf['start_date'] = pd.to_datetime(mydf['start_date'])
mydf['end_date'] = pd.to_datetime(mydf['end_date'])

# 初始化字典存储已赋值记录的id与对应日期范围
assigned = {}
for idx, row in mydf.iterrows():
    if row['id'] != '-':
        assigned[row['id']] = (row['start_date'], row['end_date'])

# 获取当前已存在的顶级id最大序号（排除A-Root）
max_top_seq = max(int(k.split('-')[1]) for k in assigned if k != 'A-Root')

# 处理未赋值的记录
for idx, row in mydf.iterrows():
    if row['id'] == '-':
        parent_id = None
        # 遍历已赋值记录，查找符合条件的父记录
        for pid, (p_start, p_end) in assigned.items():
            if row['start_date'] >= p_start and row['end_date'] <= p_end:
                parent_id = pid
                break
        
        if parent_id:
            # 统计该父记录下的子节点数量，生成子id
            child_num = sum(1 for existing_id in assigned if existing_id.startswith(f"{parent_id}."))
            new_id = f"{parent_id}.{child_num + 1}"
            mydf.loc[idx, ['id', 'parent']] = [new_id, parent_id]
            assigned[new_id] = (row['start_date'], row['end_date'])
        else:
            # 生成新的顶级id
            max_top_seq += 1
            new_id = f"A-{max_top_seq}"
            mydf.loc[idx, ['id', 'parent']] = [new_id, 'A-Root']
            assigned[new_id] = (row['start_date'], row['end_date'])

# 输出结果
print(mydf)

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