如何高效优雅地改写循环中修改DataFrame的Python代码？

问题说明 #

原代码通过循环遍历ids逐个修改DataFrame的scale列，注释明确指出这种循环修改的方式既低效又不符合Pandas的编码风格，需要更优的替代方案。

原代码 #

for id in ids:
    rows4id = df.id == id
    date2scale = get_date2scale(id)  # 原代码中date4id[id]应为笔误，示例中实际传入id
    df.loc[rows4id, "scale"] = df.loc[rows4id, "date"].map(date2scale)

示例输入 #

df = pd.DataFrame({
    "id": [0, 1, 2],
    "date": [10, 10, 10],
    "scale": [0, 0, 0]
})

ids = [1, 2]
def get_date2scale(id):
    if id == 1:
        return {10: 5.1}
    elif id == 2:
        return {10: 5.2}
    else:
        raise ValueError("Should not be called for this id")

期望输出 #

df = pd.DataFrame({
    "id": [0, 1, 2],
    "date": [10, 10, 10],
    "scale": [0, 5.1, 5.2]
})

优化方案 #

方案一：构建全局映射表批量更新 #

先整合所有需要更新的(id, date)与scale的映射关系，再通过向量化操作批量更新，避免循环修改原DataFrame。

# 构建(id, date)到scale的全局映射
id_date_scale_map = {}
for target_id in ids:
    date_scale_dict = get_date2scale(target_id)
    for date, scale in date_scale_dict.items():
        id_date_scale_map[(target_id, date)] = scale

# 仅更新匹配到的行，未匹配的保留原scale值
df["scale"] = df.apply(
    lambda row: id_date_scale_map.get((row["id"], row["date"]), row["scale"]),
    axis=1
)

方案二：使用merge合并更新 #

创建包含更新数据的辅助DataFrame，通过Pandas的merge方法完成批量更新，更贴合Pandas的向量化操作逻辑。

# 收集所有待更新的记录
update_records = []
for target_id in ids:
    date_scale_dict = get_date2scale(target_id)
    for date, scale in date_scale_dict.items():
        update_records.append({"id": target_id, "date": date, "scale": scale})

# 转换为辅助DataFrame
update_df = pd.DataFrame(update_records)

# 合并并更新scale列，保留原数据中未匹配的行
df = df.merge(update_df, on=["id", "date"], how="left", suffixes=("_original", ""))
df["scale"] = df["scale"].combine_first(df["scale_original"])
df.drop(columns=["scale_original"], inplace=True)

方案三：复合键映射+fillna #

利用id和date的复合键生成映射，通过map和fillna完成更新，代码更简洁。

# 构建复合键映射
id_date_map = {}
for target_id in ids:
    d2s = get_date2scale(target_id)
    for dt, sc in d2s.items():
        id_date_map[(target_id, dt)] = sc

# 生成复合键序列并映射，未匹配的用原scale填充
df["scale"] = df[["id", "date"]].apply(tuple, axis=1).map(id_date_map).fillna(df["scale"])

优化原因 #

• 避免循环中反复修改原DataFrame，减少Pandas内部的内存复制开销，提升运行效率。
• 符合Pandas向量化操作的编码风格，代码更简洁易读，便于后续维护和扩展。

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