如何在单条语句中为Pandas DataFrame添加两个衍生字段？

问题场景 #

我有一个Pandas DataFrame，需要根据特定条件（示例中为x > 0）添加两个计算字段SUM和DIFF，目前用两条loc语句分别实现，想知道能不能用单条语句完成，以及两种方式的性能差异。

示例代码：

import pandas as pd

df = pd.DataFrame([
    {"Name":"Axe","x":10,"y":20},
    {"Name":"Tree","x":50,"y":15},
    {"Name":"Sand","x":-10,"y":-15}
])
# 当前实现方式：两条语句分别添加字段
df.loc[df["x"] > 0, "SUM"] = df["x"] + df["y"]
df.loc[df["x"] > 0, "DIFF"] = df["x"] - df["y"]
print(df.head())

输出结果：

Name   x   y   SUM  DIFF
0    Axe  10  20  30.0 -10.0
1   Tree  50  15  65.0  35.0
2   Sand -10 -15   NaN   NaN

1. 可以用单条语句添加多个字段 #

当然可以，只需要在loc的列位置传入字段名列表，并赋值对应计算逻辑的数据集即可，这样只需要执行一次条件筛选：

方法1：向量运算直接赋值（推荐，效率最高）

先定义条件掩码，再一次性计算并赋值两列：

mask = df["x"] > 0
df.loc[mask, ["SUM", "DIFF"]] = pd.DataFrame(
    {"SUM": df.loc[mask, "x"] + df.loc[mask, "y"],
     "DIFF": df.loc[mask, "x"] - df.loc[mask, "y"]}
)

方法2：链式调用风格（assign+where）

如果习惯链式操作，可以用assign配合where实现：

df = df.assign(
    SUM=lambda d: (d["x"] + d["y"]).where(d["x"] > 0),
    DIFF=lambda d: (d["x"] - d["y"]).where(d["x"] > 0)
)

2. 两种实现方式的性能差异 #

• 小数据量场景：差异几乎可以忽略，优先考虑代码可读性即可。
• 大数据量场景：单条语句效率更高。因为两条loc语句会重复执行两次df["x"] > 0的布尔索引计算，以及两次对符合条件行的筛选操作；而单条语句只需要做一次条件判断和筛选，减少了重复计算的开销。

可以用timeit模块直观测试性能：

import timeit

def two_loc_statements():
    df = pd.DataFrame({"x": range(100000), "y": range(100000)})
    df.loc[df["x"] > 50000, "SUM"] = df["x"] + df["y"]
    df.loc[df["x"] > 50000, "DIFF"] = df["x"] - df["y"]
    return df

def single_loc_statement():
    df = pd.DataFrame({"x": range(100000), "y": range(100000)})
    mask = df["x"] > 50000
    df.loc[mask, ["SUM", "DIFF"]] = pd.DataFrame(
        {"SUM": df.loc[mask, "x"] + df.loc[mask, "y"],
         "DIFF": df.loc[mask, "x"] - df.loc[mask, "y"]}
    )
    return df

print("两条语句耗时：", timeit.timeit(two_loc_statements, number=100))
print("单条语句耗时：", timeit.timeit(single_loc_statement, number=100))

测试结果通常会显示单条语句的耗时明显低于两条语句。

内容的提问来源于stack exchange，提问作者Remis Haroon - رامز