在python中將if語句與apply結合起來

Question

蟒蛇新手。 我試圖找出基於其他列創建列的最佳方法。 理想情況下，代碼就是這樣。

df['new'] = np.where(df['Country'] == 'CA', df['x'], df['y'])

我認為這不起作用，因為它認為我正在調用整個專欄。 我試圖用 apply 做同樣的事情，但語法有問題。

df['my_col'] = df.apply(
    lambda row: 
    if row.country == 'CA':
        row.my_col == row.x
        else:
            row.my_col == row.y

我覺得必須有一個更簡單的方法。

Answer 1

這三種方法中的任何一種（ np.where 、 apply 、 mask ）似乎都有效：

df['where'] = np.where(df.country=='CA', df.x, df.y)
df['apply'] = df.apply(lambda row: row.x if row.country == 'CA' else row.y, axis=1)
mask = df.country=='CA'
df.loc[mask, 'mask'] = df.loc[mask, 'x']
df.loc[~mask, 'mask'] = df.loc[~mask, 'y']

完整的測試代碼：

import pandas as pd
import numpy as np
df = pd.DataFrame({'country':['CA','US','CA','UK','CA'], 'x':[1,2,3,4,5], 'y':[6,7,8,9,10]})
print(df)

df['where'] = np.where(df.country=='CA', df.x, df.y)
df['apply'] = df.apply(lambda row: row.x if row.country == 'CA' else row.y, axis=1)
mask = df.country=='CA'
df.loc[mask, 'mask'] = df.loc[mask, 'x']
df.loc[~mask, 'mask'] = df.loc[~mask, 'y']
print(df)

輸入：

  country  x   y
0      CA  1   6
1      US  2   7
2      CA  3   8
3      UK  4   9
4      CA  5  10

輸出

  country  x   y  where  apply  mask
0      CA  1   6      1      1   1.0
1      US  2   7      7      7   7.0
2      CA  3   8      3      3   3.0
3      UK  4   9      9      9   9.0
4      CA  5  10      5      5   5.0

Answer 2

這也可能對您有用

data = {
    'Country' : ['CA', 'NY', 'NC', 'CA'], 
    'x' : ['x_column', 'x_column', 'x_column', 'x_column'],
    'y' : ['y_column', 'y_column', 'y_column', 'y_column']
}
df = pd.DataFrame(data)
condition_list = [df['Country'] == 'CA']
choice_list = [df['x']]
df['new'] = np.select(condition_list, choice_list, df['y'])
df

您的 np.where() 看起來不錯，所以我會仔細檢查您的列是否正確標記。