Pandas按时间戳合并行并保持ID排序与值对应顺序

原始DataFrame定义 #

import pandas as pd

df = pd.DataFrame(
    {
        "IDs": [
            ["A","C"],
            ["B"],
            ["A", "B"],
            ["A","C", "D"],
            ["B", "E"],
        ],
        "values": [[1, -3], [1], [2, -2], [4, 1, 3], [2, 5]],
        "timestamp":[
            pd.to_datetime("2022-01-01"),
            pd.to_datetime("2022-01-01"),
            pd.to_datetime("2022-01-02"),
            pd.to_datetime("2022-01-03"),
            pd.to_datetime("2022-01-03"),
        ],
    },
)

原始数据输出 #

IDs         values      timestamp
0   [A, C]      [1, -3]     2022-01-01
1   [B]         [1]         2022-01-01
2   [A, B]      [2, -2]     2022-01-02
3   [A, C, D]   [4, 1, 3]   2022-01-03
4   [B, E]      [2, 5]      2022-01-03

需求说明 #

合并具有相同时间戳的行，同时满足：

• 合并后的ID列表按字母排序
• 对应的值需匹配排序后的ID位置

期望输出 #

IDs             values          timestamp
0   [A, B, C]       [1, 1, -3]      2022-01-01
1   [A, B]          [2, -2]         2022-01-02
2   [A, B, C, D, E] [4, 2, 1, 3, 5] 2022-01-03

注：实际场景中仅极小比例的行存在重复时间戳，需保证实现高效性。

高效实现方案 #

思路 #

1. 按timestamp分组，仅对存在重复时间戳的组进行处理，单一行的组直接保留
2. 每组内将ID与对应值打包为字典，自动去重（若同一时间戳同一ID出现多次，后出现的值会覆盖前一个）
3. 对字典按ID排序后，拆分回ID列表和值列表

代码实现 #

def merge_group(group):
    # 整合组内所有ID与对应值到字典
    id_value_map = {}
    for ids, vals in zip(group['IDs'], group['values']):
        id_value_map.update(dict(zip(ids, vals)))
    # 按ID排序后拆分
    sorted_pairs = sorted(id_value_map.items())
    return pd.Series({
        'IDs': [pair[0] for pair in sorted_pairs],
        'values': [pair[1] for pair in sorted_pairs]
    })

# 分组应用处理函数，保留时间戳列
result_df = df.groupby('timestamp', as_index=False).apply(merge_group)

效率说明 #

由于仅少数组存在重复时间戳，大部分组的处理只是直接返回原数据，字典整合和排序操作的开销极低，整体性能符合高效要求。

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