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 和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

性能真的很有趣，在小的DataFrames到1k更快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)')