Pandas - условно выбрать исходный столбец данных для нового столбца на основе значения строки

Есть ли функция pandas, которая позволяет выбирать из разных столбцов на основе условия? Это аналогично оператору CASE в предложении SQL Select. Например, скажем, у меня есть следующий DataFrame:

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

Я хочу выбрать из столбца 'x', когда Страна == 'США', и из столбца '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]

Ответ 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

Ответ 2

Это будет работать:

In [84]:

def func(x):
    if x['Country'] == 'USA':
        return x['x']
    if x['Country'] == 'Canada':
        return x['y']
    return NaN
foo['z'] = foo.apply(func(row), axis = 1)
foo
Out[84]:
  Country  x  y  z
0     USA  1  2  1
1  Canada  3  4  4
2  Canada  5  6  6

[3 rows x 4 columns]

Вы можете использовать loc:

In [137]:

foo.loc[foo['Country']=='Canada','z'] = foo['y']
foo.loc[foo['Country']=='USA','z'] = foo['x']
foo
Out[137]:
  Country  x  y  z
0     USA  1  2  1
1  Canada  3  4  4
2  Canada  5  6  6

[3 rows x 4 columns]

ИЗМЕНИТЬ

Несмотря на то, что громоздкий с помощью loc будет лучше масштабироваться с большими кадрами данных, поскольку здесь применяется призыв к каждой строке, в то время как использование булевской индексации будет векторизоваться.

Ответ 3

Вот общее решение для выбора произвольных столбцов с учетом значения в другом столбце.

Это имеет дополнительное преимущество для разделения логики поиска в простой структуре dict, которая упрощает ее модификацию.

import pandas as pd
df = pd.DataFrame(
    [['UK', 'burgers', 4, 5, 6],
    ['USA', 4, 7, 9, 'make'],
    ['Canada', 6, 4, 6, 'you'],
    ['France', 3, 6, 'fat', 8]],
    columns = ('Country', 'a', 'b', 'c', 'd')
)

Я перехожу к операции, в которой условный результат сохраняется во внешней структуре поиска (dict)

lookup = {'Canada': 'd', 'France': 'c', 'UK': 'a', 'USA': 'd'}

Зациклируйте pd.DataFrame для каждого столбца, хранящегося в dict, и используйте значения в таблице условий, чтобы определить, какой столбец выбрать

for k,v in lookup.iteritems():
    filt = df['Country'] == k
    df.loc[filt, 'result'] = df.loc[filt, v] # modifies in place

Чтобы дать жизненный урок

In [69]: df
Out[69]:
  Country        a  b    c     d   result
0      UK  burgers  4    5     6  burgers
1     USA        4  7    9  make     make
2  Canada        6  4    6   you      you
3  France        3  6  fat     8      fat