Pandas多层索引下按指定切片取text_len最大值行，优化嵌套循环方案

优化方案 #

1. 轻量改高效版（逻辑不变，改动最小） #

原代码的性能瓶颈主要来自两个地方：

• 每次循环重复索引df.loc[userId]，重复计算切片
• 反复调用DataFrame.append，该操作每次都会生成全新的DataFrame对象，时间复杂度为O(n²)

仅需调整结果收集逻辑和复用用户切片即可获得10倍以上的性能提升：

import pandas as pd

res_list = []
for user_id, sessions in session_dict.items():
    # 提前取出单个用户的全量数据，仅索引一次
    user_df = df.loc[user_id].reset_index(drop=True)
    for start, end in sessions:
        session_slice = user_df.iloc[start:end]
        # 取当前切片内text_len最大的行，和原逻辑一致，多最大值时取第一条
        max_row = session_slice.loc[session_slice["text_len"].idxmax()]
        res_list.append(max_row.to_dict())

# 最后一次性转为DataFrame，避免多次append的性能损耗
df_result = pd.DataFrame(res_list)

2. 向量化高性能版（适合超大数据量） #

如果你的数据量达到十万级以上，可以用全向量化操作避免Python层的嵌套循环，性能可提升数十倍：

import pandas as pd

# 给每个用户的行添加组内位置编号
df = df.reset_index()
df["group_pos"] = df.groupby("userId").cumcount()

# 将session字典转为结构化DataFrame
session_records = []
for uid, sessions in session_dict.items():
    for start, end in sessions:
        session_records.append({
            "userId": uid,
            "session_idx": len(session_records),
            "start": start,
            "end": end
        })
session_df = pd.DataFrame(session_records)

# 关联筛选出所有属于对应session区间的行
merged = df.merge(session_df, on="userId", how="inner")
merged = merged[(merged["group_pos"] >= merged["start"]) & (merged["group_pos"] < merged["end"])]

# 按session分组取text_len最大的行，清理冗余字段
df_result = merged.loc[merged.groupby("session_idx")["text_len"].idxmax()]
df_result = df_result.drop(columns=["group_pos", "session_idx", "start", "end"]).reset_index(drop=True)

注意事项 #

两种方案均和原代码逻辑完全对齐，同一session内存在多个相同text_len最大值时，默认取排序靠前的第一条，如需保留所有最大值可将idxmax替换为自定义的等值筛选逻辑。

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