如何在Pandas中构建以DataFrame行为键的a-b行配对字典？

背景 #

• DataFrame每行代表一个事件，包含Seat标识符和Type字段（值为'a'或'b'）
• 需要将每个'a'条目与同一Seat下后续出现的'b'条目配对
• DataFrame已按Seat和时间排序，确保'a'条目一定在对应'b'条目之前

问题 #

我想用字典存储这些配对，理想结构是{a_entry: b_entry}，但DataFrame行是不可哈希的，没法直接当字典键。我偏好字典是因为可以像下面这样轻松合并两行数据：

merged_entries = []

for a_entry, b_entry in pairs_dict.items():
    code = '19' if a_entry['Asleep'] and b_entry['Asleep'] else '11'
    if code == '19':
        logging.info(f"Changed code from 11 to 19 for seat {a_entry['Seat']}")
    
    merged_entry = {
        'Seat': a_entry['Seat'],
        'Starttime': a_entry['Starttime'],
        'Endtime': b_entry['Endtime'],
        'Duration': b_entry['Endtime'] - a_entry['Starttime'],
        'Code': code,
        'Notes': None,
        'Asleep': a_entry['Asleep'],
        'Concert': a_entry['Concert'],
        'Camera_Position': a_entry['Camera_Position'],
        'Composer': a_entry['Composer'],
        'Movement': a_entry['Movement'],
        'File_Name': a_entry['File_Name']
    }
    merged_entries.append(merged_entry)

示例数据 #

from datetime import datetime
import pandas as pd

# 示例数据
data = {
    'Seat': [1, 1, 1, 2, 2, 3, 3, 3, 3],
    'Starttime': [datetime(2023, 1, 1, 9, 0),
                  datetime(2023, 1, 1, 9, 30),
                  datetime(2023, 1, 1, 10, 0),
                  datetime(2023, 1, 1, 9, 15),
                  datetime(2023, 1, 1, 9, 45),
                  datetime(2023, 1, 1, 10, 10),
                  datetime(2023, 1, 1, 10, 30),
                  datetime(2023, 1, 1, 10, 50),
                  datetime(2023, 1, 1, 11, 10)],
    'Endtime': [datetime(2023, 1, 1, 9, 20),
                datetime(2023, 1, 1, 9, 50),
                datetime(2023, 1, 1, 10, 20),
                datetime(2023, 1, 1, 9, 35),
                datetime(2023, 1, 1, 10, 5),
                datetime(2023, 1, 1, 10, 20),
                datetime(2023, 1, 1, 10, 50),
                datetime(2023, 1, 1, 11, 0),
                datetime(2023, 1, 1, 11, 20)],
    'Type': ['a', 'b', 'b', 'a', 'b', 'a', 'a', 'b', 'b'],
    'Asleep': [False, True, False, True, True, False, True, True, False],
    'Code': [11, 11, 11, 11, 11, 11, 11, 11, 11],
    'Notes': [None, None, None, None, None, None, None, None, None],
    'Concert': ['Concert1', 'Concert1', 'Concert1', 'Concert2', 'Concert2', 'Concert3', 'Concert3', 'Concert3', 'Concert3'],
    'Camera_Position': ['Front', 'Front', 'Front', 'Back', 'Back', 'Left', 'Left', 'Left', 'Left'],
    'Composer': ['Bach', 'Bach', 'Bach', 'Mozart', 'Mozart', 'Beethoven', 'Beethoven', 'Beethoven', 'Beethoven'],
    'Movement': ['Adagio', 'Adagio', 'Adagio', 'Allegro', 'Allegro', 'Fur Elise', 'Fur Elise', 'Fur Elise', 'Fur Elise'],
    'File_Name': ['file1.csv', 'file1.csv', 'file1.csv', 'file2.csv', 'file2.csv', 'file3.csv', 'file3.csv', 'file3.csv', 'file3.csv']
}
df = pd.DataFrame(data)

已尝试方案 #

• 将行转换为元组或字符串，但这样会失去以DataFrame形式操作数据的便利性
• 使用列表等其他数据结构，但更倾向于字典的O(1)检索效率

解决方案 #

方案1：用行索引作为字典键 #

利用DataFrame的唯一可哈希索引作为键，存储对应的a、b行：

import logging

# 按Seat分组，处理每个座位的a/b配对
pairs_dict = {}
# 记录每个Seat待配对的a行索引
pending_a = {}

for idx, row in df.iterrows():
    seat = row['Seat']
    if row['Type'] == 'a':
        # 记录当前a行的索引，等待配对后续的b行
        pending_a[seat] = idx
    elif row['Type'] == 'b' and seat in pending_a:
        # 找到配对，将a行索引作为键，b行作为值存入字典
        a_idx = pending_a.pop(seat)
        pairs_dict[a_idx] = row.to_dict()

# 合并条目
merged_entries = []
for a_idx, b_row in pairs_dict.items():
    a_row = df.loc[a_idx]
    code = '19' if a_row['Asleep'] and b_row['Asleep'] else '11'
    if code == '19':
        logging.info(f"Changed code from 11 to 19 for seat {a_row['Seat']}")
    
    merged_entry = {
        'Seat': a_row['Seat'],
        'Starttime': a_row['Starttime'],
        'Endtime': b_row['Endtime'],
        'Duration': b_row['Endtime'] - a_row['Starttime'],
        'Code': code,
        'Notes': None,
        'Asleep': a_row['Asleep'],
        'Concert': a_row['Concert'],
        'Camera_Position': a_row['Camera_Position'],
        'Composer': a_row['Composer'],
        'Movement': a_row['Movement'],
        'File_Name': a_row['File_Name']
    }
    merged_entries.append(merged_entry)

# 转为DataFrame查看结果
merged_df = pd.DataFrame(merged_entries)

方案2：分组后直接配对生成字典 #

通过分组和排序，直接将每个a行与对应的b行配对：

# 按Seat分组，分别提取a和b行
grouped = df.groupby('Seat')
pairs_dict = {}

for seat, group in grouped:
    a_rows = group[group['Type'] == 'a']
    b_rows = group[group['Type'] == 'b']
    # 按时间顺序一一配对（因为已排序）
    for a, b in zip(a_rows.itertuples(), b_rows.itertuples()):
        # 用a行的索引作为键，b行转为字典
        pairs_dict[a.Index] = b._asdict()

# 后续合并逻辑同方案1

方案3：生成唯一标识符作为键 #

如果索引可能重复，可以为每个a行生成唯一标识符（比如结合Seat和Starttime）作为键：

pairs_dict = {}
pending_a = {}

for idx, row in df.iterrows():
    seat = row['Seat']
    if row['Type'] == 'a':
        # 生成唯一键：Seat+Starttime的字符串形式
        unique_key = f"{seat}_{row['Starttime'].isoformat()}"
        pending_a[unique_key] = row.to_dict()
    elif row['Type'] == 'b' and pending_a:
        # 找到当前Seat对应的最后一个待配对a行
        matching_keys = [k for k in pending_a.keys() if k.startswith(f"{seat}_")]
        if matching_keys:
            # 按时间取最新的a行（因为已排序，最后一个就是最新的）
            matching_keys.sort()
            a_key = matching_keys[-1]
            a_row = pending_a.pop(a_key)
            pairs_dict[a_key] = row.to_dict()

这些方案既保留了字典的高效检索特性，又能正常操作DataFrame行的数据。

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