Pandas计算连通pd.Intervals区间ID及连续区间长度

• 按固定长度5对时间轴划分左闭右开分箱，例如[0, 5)、[5, 10)依次类推
• 现有timestamp列存储大于等于0的时间戳，可通过pd.cut()匹配每个时间戳对应的分箱，例如时间戳3.0对应分箱[0,5)
• 无对应时间戳的分箱不会出现在time_bin列中，因此列中可能存在区间缺口：例如列中同时出现[5,10)、[15,20)时，代表缺失[10,15)区间，已知timestamp列已完成排序

初始DataFrame包含group_id、timestamp、time_bin三列，需新增两个计算列：

1. connected_interval：标记当前行所属连通区间的整数ID，连通定义为相邻分箱无缺口，例如[0,5)、[5,10)、[10,15)属于同一连通块，分配相同ID，不同连通块ID从0开始按出现顺序递增
2. conn_interv_len：标记当前行所属最大连通区间的总长度，例如上述三个连续分箱组成的连通区间总长度为15，连通块内所有行的该列值统一为15

测试数据构造代码 #

import pandas as pd
df = pd.DataFrame({
    "group_id":['A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B'],
    "timestamp": [0.0, 3.0, 9.0, 24.2, 30.2, 0.0, 136.51, 222.0, 237.0, 252.0],
    "time_bin": [
        pd.Interval(0, 5, closed='left'), pd.Interval(0, 5, closed='left'),
        pd.Interval(5, 10, closed='left'), pd.Interval(20, 25, closed='left'),
        pd.Interval(30, 35, closed='left'), pd.Interval(0, 5, closed='left'),
        pd.Interval(135, 140, closed='left'), pd.Interval(220, 225, closed='left'),
        pd.Interval(235, 240, closed='left'), pd.Interval(250, 255, closed='left')
    ]
})

预期输出结果 #

按group_id分组逐组计算，逻辑如下：

• 对每组数据提取按出现顺序排列的唯一分箱，判断相邻分箱是否连通（前一个分箱右边界等于后一个分箱左边界即为连通）
• 在分箱缺口处做累加标记，生成每个唯一分箱对应的连通区间ID
• 统计每个连通ID包含的唯一分箱数量，乘以单分箱长度5得到连通区间总长度
• 将分箱对应的ID、长度映射回原表所有行即可

可直接运行的实现代码 #

BIN_STEP = 5

def calc_conn_metrics(sub_df: pd.DataFrame) -> pd.DataFrame:
    # 提取按原顺序排列的唯一分箱
    unique_bins = sub_df['time_bin'].drop_duplicates().reset_index(drop=True)
    # 识别缺口位置，生成连通区间ID
    is_gap = unique_bins.shift(1).apply(lambda x: x.right if pd.notnull(x) else None) != unique_bins.apply(lambda x: x.left)
    conn_id = is_gap.cumsum() - 1
    # 计算每个连通区间的总长度
    conn_len = conn_id.value_counts().mul(BIN_STEP).to_dict()
    # 构建分箱到属性的映射字典
    bin_to_id = dict(zip(unique_bins, conn_id))
    bin_to_len = {b: conn_len[bin_to_id[b]] for b in unique_bins}
    # 映射回当前分组的所有行
    sub_df['connected_interval'] = sub_df['time_bin'].map(bin_to_id)
    sub_df['conn_interv_len'] = sub_df['time_bin'].map(bin_to_len)
    return sub_df

# 分组应用计算，重置索引保证结果顺序
df = df.groupby('group_id', group_keys=False).apply(calc_conn_metrics).reset_index(drop=True)

运行代码后输出结果与预期完全一致。

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