Pandas多条件赋值：np.where、loc、apply详解

2025-03-03 08:23:01

多条件赋值：Pandas 数据处理技巧

处理数据时，根据多个条件给已有列赋值挺常见的。这次碰到的问题是这样：有一个 DataFrame，包含 "gender"、"pet1" 和 "pet2" 三列，要根据这三列的值来确定一个新列 "points" 的值。

具体规则如下：

a. 如果 "gender" 是 "male"，并且 "pet1" 和 "pet2" 相等，那么 "points" 等于 5。
b. 如果 "gender" 是 "female"，并且 "pet1" 是 "cat" 或者 "dog"，那么 "points" 等于 5。
c. 其他所有情况，"points" 等于 0。

原始数据长这样：

    gender    pet1      pet2
0   male      dog       dog
1   male      cat       cat
2   male      dog       cat
3   female    cat       squirrel
4   female    dog       dog
5   female    squirrel  cat
6   squirrel  dog       cat

目标结果如下：

    gender    pet1      pet2      points
0   male      dog       dog       5
1   male      cat       cat       5
2   male      dog       cat       0
3   female    cat       squirrel  5
4   female    dog       dog       5
5   female    squirrel  cat       0
6   squirrel  dog       cat       0

下面来看看怎么用 Pandas 解决这个问题。

问题分析

这个问题的核心在于如何将多个条件组合起来，并根据这些条件判断来设置新列的值。Pandas 提供了多种方法来实现这种条件赋值，接下来我们将逐一介绍这些方法，并分析它们的优缺点。

解决方案

1. 使用 `np.where()`

np.where() 是 NumPy 库中的一个函数，可以根据条件返回不同的值。在 Pandas 中，可以结合多个条件使用 np.where() 来实现多条件赋值。

原理：

np.where() 函数接受三个参数：条件、满足条件时的值、不满足条件时的值。我们可以嵌套多个 np.where() 来处理多个条件。

代码示例：

import pandas as pd
import numpy as np

data = {'gender': ['male', 'male', 'male', 'female', 'female', 'female', 'squirrel'],
        'pet1': ['dog', 'cat', 'dog', 'cat', 'dog', 'squirrel', 'dog'],
        'pet2': ['dog', 'cat', 'cat', 'squirrel', 'dog', 'cat', 'cat']}
df = pd.DataFrame(data)

df['points'] = np.where((df['gender'] == 'male') & (df['pet1'] == df['pet2']), 5,
                       np.where((df['gender'] == 'female') & (df['pet1'].isin(['cat', 'dog'])), 5, 0))

print(df)

安全建议：

np.where 本身没有安全风险。但在使用过程中，需确保条件的逻辑正确，否则可能产生意料之外的结果。

进阶使用技巧：

当条件很多，可以把条件拆分保存，增加可读性。

condition1 = (df['gender'] == 'male') & (df['pet1'] == df['pet2'])
condition2 = (df['gender'] == 'female') & (df['pet1'].isin(['cat', 'dog']))
df['points'] = np.where(condition1 | condition2, 5, 0)

2. 使用 `pandas.DataFrame.loc[]`

loc[] 是 Pandas 中用于基于标签进行索引和选择数据的方法。它可以用来设置满足特定条件的数据的值。

原理：

loc[] 可以接受一个布尔数组作为索引，选择 DataFrame 中满足条件的行，然后对这些行进行赋值。

代码示例：

import pandas as pd

data = {'gender': ['male', 'male', 'male', 'female', 'female', 'female', 'squirrel'],
        'pet1': ['dog', 'cat', 'dog', 'cat', 'dog', 'squirrel', 'dog'],
        'pet2': ['dog', 'cat', 'cat', 'squirrel', 'dog', 'cat', 'cat']}
df = pd.DataFrame(data)

df['points'] = 0  # 先全部赋值为0
df.loc[(df['gender'] == 'male') & (df['pet1'] == df['pet2']), 'points'] = 5
df.loc[(df['gender'] == 'female') & (df['pet1'].isin(['cat', 'dog'])), 'points'] = 5

print(df)

安全建议：

loc[] 没有安全风险，注意条件的正确性。

进阶使用技巧：

可以用.copy()避免 SettingWithCopyWarning

某些 Pandas 操作可能会返回数据的视图（view）而不是副本（copy）。当尝试在视图上进行赋值时，Pandas 可能会发出 SettingWithCopyWarning 警告。

可以使用链式索引，但需要注意避免 SettingWithCopyWarning：

df = pd.DataFrame(data).copy() # 建立副本

3. 使用`apply()`方法

apply() 方法可以将一个自定义函数应用到 DataFrame 的每一行或每一列。

原理：

我们可以定义一个函数，该函数接受一行数据作为输入，并根据条件返回 "points" 列的值。然后，使用 apply() 方法将这个函数应用到 DataFrame 的每一行。

代码示例：

import pandas as pd

data = {'gender': ['male', 'male', 'male', 'female', 'female', 'female', 'squirrel'],
        'pet1': ['dog', 'cat', 'dog', 'cat', 'dog', 'squirrel', 'dog'],
        'pet2': ['dog', 'cat', 'cat', 'squirrel', 'dog', 'cat', 'cat']}
df = pd.DataFrame(data)

def calculate_points(row):
    if row['gender'] == 'male' and row['pet1'] == row['pet2']:
        return 5
    elif row['gender'] == 'female' and row['pet1'] in ['cat', 'dog']:
        return 5
    else:
        return 0

df['points'] = df.apply(calculate_points, axis=1)

print(df)

安全建议：

因为apply()使用自定义函数,要保证函数的逻辑没问题。如果自定义函数逻辑复杂或者引用外部变量，需要仔细检查。

进阶使用技巧:

可以和 lambda 函数结合使用。不过因为需要写多个判断，会显得可读性比较差。
对于非常复杂的逻辑，定义命名函数，然后在 apply() 中调用命名函数。

df['points'] = df.apply(lambda row: 5 if (row['gender'] == 'male' and row['pet1'] == row['pet2']) or \
                                        (row['gender'] == 'female' and row['pet1'] in ['cat', 'dog']) else 0, axis=1)

4. 条件选择 + `mask()`/`where()`

这种方法结合了条件选择和 Pandas 的 mask() 或 where() 方法。

原理：

先创建一个布尔掩码（mask），表示满足条件的行。然后，使用 mask() 或 where() 方法将不满足条件的行的值替换为 0，满足的则替换为 5。mask() 方法在条件为 True 时替换值，where() 方法在条件为 False 时替换值.

代码示例：

import pandas as pd

data = {'gender': ['male', 'male', 'male', 'female', 'female', 'female', 'squirrel'],
        'pet1': ['dog', 'cat', 'dog', 'cat', 'dog', 'squirrel', 'dog'],
        'pet2': ['dog', 'cat', 'cat', 'squirrel', 'dog', 'cat', 'cat']}
df = pd.DataFrame(data)

condition1 = (df['gender'] == 'male') & (df['pet1'] == df['pet2'])
condition2 = (df['gender'] == 'female') & (df['pet1'].isin(['cat', 'dog']))

df['points'] = (condition1 | condition2).mask(cond=(condition1 | condition2), other=5).fillna(0).astype(int) #或者.where(~(condition1 | condition2), 5)

# df['points']=(condition1 | condition2) 这一步会产生布尔值 True,False, .astype(int) 会把它们转换为 1，0
# 所以上一步可以简化
# df['points'] = (condition1 | condition2).astype(int) * 5
print(df)

安全建议

没有直接安全问题，但仍需小心条件和数据类型问题.

进阶使用技巧:

可以用.clip()设置最大/最小值, 结合条件可以灵活地给某一列数值设置范围。

df['points'] = (condition1 | condition2).astype(int)
df['points'] = df['points'].clip(lower=0, upper=5) # 把结果限制到0-5之间, 小于0的变为0， 大于5的变为5. 本例中和直接 *5 等价。