Панды: как сравнить несколько ячеек со списком / кортежем

Question

Мне нужно сравнить несколько столбцов в кадре данных в целом, например:

df = pd.DataFrame({'A':[1,1,3],'B':[4,5,6]})

#Select condition: If df['A'] == 1 and df['B'] == 4, then pick up this row. 

Для этого простого примера я могу использовать метод ниже:

df.loc[(df['A']==1)&(df['B']==4),'A':'B']

Однако на самом деле мой фрейм данных имеет десятки столбцов, которые следует сравнивать как единое целое. Приведенное выше решение будет очень и очень грязным, если я решу перечислить все из них. Поэтому я думаю, что если рассматривать их как единое целое, то сравнение со списком может решить проблему:

#something just like this:
df.loc[df.loc[:,'A':'B']==[1,4],'A':'B')]

Не сработало. Поэтому мне пришла в голову идея, что сначала нужно объединить все нужные столбцы в новый столбец в качестве значения списка, а затем сравнить этот новый столбец со списком. Последнее было решено в Pandas: сравните список объектов в Series

Несмотря на то, что в общем случае я раскрыл свое дело, я все еще хочу знать, есть ли более простой способ решить эту проблему? Спасибо.

Answer 1

Или используйте [[]] для получения нескольких столбцов:

df[(df[['A','B']].values==[1,4]).all(1)]

Демо-версия:

>>> df = pd.DataFrame({'A':[1,1,3],'B':[4,5,6]})
>>> df[(df[['A','B']].values==[1,4]).all(1)]
   A  B
0  1  4
>>> 

Answer 2

Вы можете использовать логическую маску через представление ваших данных в массиве NumPy:

df = pd.DataFrame({'A':[1,1,3],'B':[4,5,6]})

res = df[(df.loc[:, 'A':'B'].values == [1, 4]).all(1)]

print(res)

   A  B
0  1  4

В этой ситуации никогда не объединяет ваши столбцы в одну серию списков. Это неэффективно, поскольку вы потеряете все преимущества векторизации, и любая последующая обработка будет включать циклы уровня Python.