Как создать маску для датафрейма по списку подходящих значений в колонке?

Question

Есть датафрейм по продажам автомобилей. Каждая строка соответствует событию продажи автомобиля. Одна из колонок - "dealers". В ней указаны дилер, осуществивший продажу. Я хочу сделать маску для датафрейма, которая оставит только строки, соответствующие дилерам, указанным в списке. Как правильно это сделать? Как правильно сформировать маску? Пытался через in:

mask = sales['dealers'] in list_of_dealers

Выводит ошибку:

ValueError: The truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all().

Answer 1

Один из подходов отбора строк по заданному перечню значений - использовать метод Series.isin, который вернет результат проверки item in given_list для каждого элемента последовательности:

import pandas as pd

sales = pd.DataFrame({
    'dealers': [*'ABCDABCD'],
    'foo': [*'xyzxwzyy'],
    'bar': [*'12642136']
})

list_of_dealers = [*'BD']
mask = sales['dealers'].isin(list_of_dealers)
print(mask)

# 0    False
# 1     True
# 2    False
# 3     True
# 4    False
# 5     True
# 6    False
# 7     True
# Name: dealers, dtype: bool

Если нужен обратный эффект - все, кроме указанных продавцов, - дополнительно применяем оператор поэлементного инвертирования ~:

mask = ~sales['dealers'].isin(list_of_dealers)

Чуть менее очевидный, но вполне рабочий способ отбора строк - использовать индексы в паре со свойством DataFrame.loc:

selected = (
    sales
    .set_index('dealers')
    .loc[list_of_dealers]
    .reset_index()
)
print(selected)

#   dealers foo bar
# 0       B   y   2
# 1       B   z   1
# 2       D   x   4
# 3       D   y   6

Для обратного эффекта - все, кроме заданных продавцов - применяем метод DataFrame.drop вместо свойства loc:

print(
    sales
    .set_index('dealers')
    .drop(list_of_dealers)
    .reset_index()
)

#   dealers foo bar
# 0       A   x   1
# 1       C   z   6
# 2       A   w   2
# 3       C   y   3

БЛОГ НА HUSL

Как создать маску для датафрейма по списку подходящих значений в колонке?

Ответы (1 шт):