pandas 使用条件 lambda 表达式分配多个列

Question

我想使用 .assign() 方法将 2 列（cat_a、cat_b）添加到 DataFrame df。但是我没有让代码正常工作...

import pandas as pd
np.random.seed(999)
num = 10
df = pd.DataFrame({'id': np.random.choice(range(1000, 10000), num, replace=False),
                   'sex': np.random.choice(list('MF'), num, replace=True),
                   'year': np.random.randint(1980, 1990, num)})
print(df)

     id sex  year
0  3461   F  1983
1  8663   M  1988
2  6615   M  1986
3  5336   M  1982
4  3756   F  1984
5  8653   F  1989
6  9362   M  1985
7  3944   M  1981
8  3334   F  1986
9  6135   F  1988

这应该是新列 cat_a 和 cat_b

的值

# cat_a
list(map(lambda y: 'A' if y <= 1985 else 'B', df.year))
['A', 'B', 'B', 'A', 'A', 'B', 'A', 'A', 'B', 'B']

# cat_b
list(map(lambda s, y: 1 if s == 'M' and y <= 1985 else (2 if s == 'M' else (3 if y < 1985 else 4)), df.sex, df.year))
[3, 2, 2, 1, 3, 4, 1, 1, 4, 4]

正在尝试 .assign() 方法的语法：

df.assign(cat_a = 'AB', cat_b = 1234)
print(df)

     id sex  year cat_a  cat_b
0  3461   F  1983    AB   1234
1  8663   M  1988    AB   1234
2  6615   M  1986    AB   1234
3  5336   M  1982    AB   1234
4  3756   F  1984    AB   1234
5  8653   F  1989    AB   1234
6  9362   M  1985    AB   1234
7  3944   M  1981    AB   1234
8  3334   F  1986    AB   1234
9  6135   F  1988    AB   1234

替换虚拟值会出现错误：

df.assign(cat_a = lambda x: 'A' if x.year <= 1985 else 'B',
          cat_b = lambda x: 1 if x.sex == 'M' and x.year <= 1985 
                              else (2 if x.sex == 'M'
                                      else (3 if x.year < 1985
                                              else 4
                                           )
                                   )
         )

非常欢迎任何有关如何使代码正常工作的建议！
我有解决方法，但我想使用 .assign() 方法获得结果。

Answer 1

使用矢量化解决方案 numpy.where and numpy.select:

m1 = df.year <= 1985
m2 = df.sex == 'M'

a = np.where(m1, 'A', 'B')
b = np.select([m1 & m2, ~m1 & m2, m1 & ~m2], [1,2,3], default=4)

df = df.assign(cat_a = a, cat_b = b)
print (df)
     id sex  year cat_a  cat_b
0  3461   F  1983     A      3
1  8663   M  1988     B      2
2  6615   M  1986     B      2
3  5336   M  1982     A      1
4  3756   F  1984     A      3
5  8653   F  1989     B      4
6  9362   M  1985     A      1
7  3944   M  1981     A      1
8  3334   F  1986     B      4
9  6135   F  1988     B      4

验证:

a = list(map(lambda y: 'A' if y <= 1985 else 'B', df.year))
b = list(map(lambda s, y: 1 if s == 'M' and y <= 1985 else (2 if s == 'M' else (3 if y < 1985 else 4)), df.sex, df.year))

df = df.assign(cat_a = a, cat_b = b)
print (df)
     id sex  year cat_a  cat_b
0  3461   F  1983     A      3
1  8663   M  1988     B      2
2  6615   M  1986     B      2
3  5336   M  1982     A      1
4  3756   F  1984     A      3
5  8653   F  1989     B      4
6  9362   M  1985     A      1
7  3944   M  1981     A      1
8  3334   F  1986     B      4
9  6135   F  1988     B      4

性能真的很有趣，在 1k 的小 DataFrames 中更快 mapping，对于更大的 DataFrames 更好地使用 numpy 解决方案：

np.random.seed(999)

def mapping(df):
    a = list(map(lambda y: 'A' if y <= 1985 else 'B', df.year))
    b = list(map(lambda s, y: 1 if s == 'M' and y <= 1985 else (2 if s == 'M' else (3 if y < 1985 else 4)), df.sex, df.year))

    return df.assign(cat_a = a, cat_b = b)

def vec(df):
    m1 = df.year <= 1985
    m2 = df.sex == 'M'
    a = np.where(m1, 'A', 'B')
    b = np.select([m1 & m2, ~m1 & m2, m1 & ~m2], [1,2,3], default=4)
    return df.assign(cat_a = a, cat_b = b)

def make_df(n):
    df = pd.DataFrame({'id': np.random.choice(range(10, 1000000), n, replace=False),
                   'sex': np.random.choice(list('MF'), n, replace=True),
                   'year': np.random.randint(1980, 1990, n)})
    return df

perfplot.show(
    setup=make_df,
    kernels=[mapping, vec],
    n_range=[2**k for k in range(2, 18)],
    logx=True,
    logy=True,
    equality_check=False,  # rows may appear in different order
    xlabel='len(df)')

pandas 使用条件 lambda 表达式分配多个列

pandas assign multiple columns with conditional lambda expression

lambda

multiple-columns

python-3.x

assign

pandas