Как выбрать строки из DataFrame на основе значений столбцов?

Вот простой пример, который иллюстрирует, как использовать маску в pandas для фильтрации данных в DataFrame.

from pandas import DataFrame

# Создаем набор данных
d = {'Revenue': [100, 111, 222], 
     'Cost': [333, 444, 555]}
df = DataFrame(d)

# Создаем маску: возвращает True, когда значение в колонке "Revenue" равно 111
mask = df['Revenue'] == 111

print(mask)

# Результат:
# 0    False
# 1     True
# 2    False
# Name: Revenue, dtype: bool

# Выбираем строки из df, где Revenue = 111
filtered_df = df[mask]

# Результат:
#    Cost  Revenue
# 1   444      111

В этом коде мы сначала создаем DataFrame с данными о доходах и затратах. Затем мы создаем маску, которая проверяет, равно ли значение в столбце "Revenue" 111. После этого мы применяем эту маску для фильтрации DataFrame, чтобы получить строки, где доход составляет 111. Вывод показывает только ту строку, где это условие выполнено.

Чтобы добавить: Вы также можете использовать df.groupby('column_name').get_group('column_desired_value').reset_index(), чтобы создать новый DataFrame, в котором указанная колонка имеет определенное значение. Например:

import pandas as pd
df = pd.DataFrame({'A': 'foo bar foo bar foo bar foo foo'.split(),
                   'B': 'one one two three two two one three'.split()})
print("Исходный DataFrame:")
print(df)

b_is_two_dataframe = pd.DataFrame(df.groupby('B').get_group('two').reset_index()).drop('index', axis=1)
# ПРИМЕЧАНИЕ: финальный drop нужен для удаления лишней колонки индекса, возвращаемой объектом groupby
print('Подмножество DataFrame, где B равно two:')
print(b_is_two_dataframe)

При выполнении этого кода вы увидите следующий результат:

Исходный DataFrame:
     A      B
0  foo    one
1  bar    one
2  foo    two
3  bar  three
4  foo    два
5  bar    два
6  foo    один
7  foo  три
Подмножество DataFrame, где B равно two:
     A    B
0  foo  два
1  foo  два
2  bar  два

Вы также можете использовать метод .apply:

df.apply(lambda row: row[df['B'].isin(['one', 'three'])])

Это на самом деле работает по строкам (то есть функция применяется к каждой строке).

Результат будет следующим:

   A      B  C   D
0  foo    one  0   0
1  bar    one  1   2
3  bar  three  3   6
6  foo    one  6  12
7  foo  three  7  14

Результат такой же, как и при использовании метода, упомянутого пользователем @unutbu:

df[[df['B'].isin(['one', 'three'])]]

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

Для этого воспользуйтесь следующим кодом:

my_df = df.set_index(column_name)
my_dict = my_df.to_dict('index')

После создания словаря my_dict вы сможете выполнять запросы следующим образом:

if some_value in my_dict.keys():
   my_result = my_dict[some_value]

Обратите внимание, что если в column_name есть повторяющиеся значения, вы не сможете создать словарь. В этом случае, вы можете использовать следующий код для получения результата:

my_result = my_df.loc[some_value]

Таким образом, в зависимости от наличия дубликатов в вашем столбце, вы выберете подходящий способ для поиска данных.

Вы можете использовать loc (квадратные скобки) с функцией:

# Series
s = pd.Series([1, 2, 3, 4]) 
s.loc[lambda x: x > 1]
# s[lambda x: x > 1]

Вывод:

1    2
2    3
3    4
dtype: int64

или

# DataFrame
df = pd.DataFrame({'A': [1, 2, 3], 'B': [10, 20, 30]})
df[lambda x: (x['A'] != 1) & (x['B'] != 30)]

Вывод:

   A   B
1  2  20

Таким образом, вы можете использовать лямбда-функции для фильтрации данных в Series и DataFrame.