Используя условие, выберите нужные столбцы в панде DataFrame

Question

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

df = pd.DataFrame(
    [['Y', 'Cat', 'no', 'yes', 6],
    ['Y', 4, 7, 9, 'dog'],
    ['N', 6, 4, 6, 'pig'],
    ['N', 3, 6, 'beer', 8]],
    columns = ('Data', 'a', 'b', 'c', 'd')
)

Мое условие, которое не работает:

if (df['Data']=='Y') & (df['Data']=='N'):
    df3=df.loc[:,['Data', 'a', 'b', 'c']]
else:
    df3=df.loc[:,['Data', 'a', 'b']]

Я хочу, чтобы новая таблица содержала данные, соответствующие следующим критериям:

Если df.Data имеет значения 'Y' и 'N', новая таблица получит столбцы ('Data', 'a', 'b')

Если нет, новая таблица получает столбцы ('Data', 'a', 'b', 'c')

 Data    a   b
0    Y  Cat  no
1    Y    4   7
2    N    6   4
3    N    3   6

  Data    a   b     c
0    Y  Cat  no   yes
1    Y    4   7     9
2    Y    6   4     6
3    Y    3   6  beer

Answer 1

if all(df.Data.unique() == ['Y','N']) == True:
    df3 = df[['Data', 'a', 'b', 'c']]
else:
    df3 = df[['Data','a','b']]

Answer 2

Вы сравниваете серию с символом, а не проверяете существование одного логического результата.Вместо этого вы можете использовать pd.Series.any, который возвращает True, если какое-либо значение в серии равно True:

if (df['Data']=='Y').any() & (df['Data']=='N').any():
    # do something

Альтернативный метод - использовать pd.DataFrame.drop с троичным оператором:

df = df.drop(['d'] if set(df['Data']) == {'Y', 'N'} else ['c', 'd'], 1)

print(df)

  Data    a   b     c
0    Y  Cat  no   yes
1    Y    4   7     9
2    N    6   4     6
3    N    3   6  beer