Pandas根据行条件转换DataFrame并提取对应列名的实现方法

实现思路 #

• 第一步筛选所有speed_t开头的速度维度列，排除car id等非速度列
• 第二步对每一行的数据，按100的阈值筛选符合条件的列名，拼接为逗号分隔的字符串
• 第三步将car id列与两个拼接得到的汇总列合并，得到最终输出

完整代码实现 #

示例数据构造 #

import pandas as pd

# 构造题目给出的输入DataFrame
df = pd.DataFrame({
    "car id": ["1.", "2.", "3."],
    "speed_t1": [90, 110, 40],
    "speed_t2": [80, 130, 110],
    "speed_t3": [120, 140, 20],
    "speed_t4": [34, 99, 110]
})

核心转换逻辑 #

小数据量直接用行apply实现，代码简单易读：

# 提取所有速度列
speed_cols = [col for col in df.columns if col.startswith("speed_t")]

# 按行处理生成两个汇总列
result = df[["car id"]].copy()
result["normal_speed_t"] = df.apply(
    lambda x: ",".join([c for c in speed_cols if x[c] <= 100]), axis=1
)
result["high_speed_t"] = df.apply(
    lambda x: ",".join([c for c in speed_cols if x[c] > 100]), axis=1
)

大数据量优化实现 #

如果数据量超过10万行，推荐用向量化掩码操作提升运行效率：

speed_cols = [col for col in df.columns if col.startswith("speed_t")]
mask_normal = df[speed_cols] <= 100
mask_high = df[speed_cols] > 100

result = df[["car id"]].copy()
result["normal_speed_t"] = mask_normal.apply(lambda x: ",".join(speed_cols[x]), axis=1)
result["high_speed_t"] = mask_high.apply(lambda x: ",".join(speed_cols[x]), axis=1)

运行后输出的result和题目给出的预期输出完全一致。

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