python - Pandas - 根据行值有条件地为新列选择数据源列

Question

是否有一个 pandas 函数可以根据条件从不同的列中进行选择？这类似于 SQL Select 子句中的 CASE 语句。例如，假设我有以下 DataFrame:

foo = DataFrame(
    [['USA',1,2],
    ['Canada',3,4],
    ['Canada',5,6]], 
    columns = ('Country', 'x', 'y')
)

我想在 Country=='USA' 时从列 'x' 中选择，在 Country=='Canada' 时从列'y' 中选择，结果如下所示:

  Country  x  y  z
0     USA  1  2  1
1  Canada  3  4  4
2  Canada  5  6  6

[3 rows x 4 columns]

Answer 1

使用 DataFrame.where的 other 参数和 pandas.concat :

>>> import pandas as pd
>>>
>>> foo = pd.DataFrame([
...     ['USA',1,2],
...     ['Canada',3,4],
...     ['Canada',5,6]
... ], columns=('Country', 'x', 'y'))
>>>
>>> z = foo['x'].where(foo['Country'] == 'USA', foo['y'])
>>> pd.concat([foo['Country'], z], axis=1)
  Country  x
0     USA  1
1  Canada  4
2  Canada  6

如果你想要z作为列名，指定keys:

>>> pd.concat([foo['Country'], z], keys=['Country', 'z'], axis=1)
  Country  z
0     USA  1
1  Canada  4
2  Canada  6