Фильтрация строк DataFrame по наличию значения в колонке из заданного списка значений

Вы можете использовать метод isin для фильтрации данных в DataFrame. Например, если у вас есть DataFrame rpt и список идентификаторов акций stk_list, вы можете отфильтровать строки, где столбец STK_ID содержится в этом списке, следующим образом:

filtered_rpt = rpt[rpt['STK_ID'].isin(stk_list)]

Таким образом, в переменной filtered_rpt будут храниться только те строки из DataFrame rpt, которые соответствуют идентификаторам акций из stk_list.

Вы также можете добиться аналогичных результатов, использовав метод query и знак @:

Например:

import pandas as pd

df = pd.DataFrame({'A': [1, 2, 3], 'B': ['a', 'b', 'f']})
df = pd.DataFrame({'A': [5, 6, 3, 4], 'B': [1, 2, 3, 5]})
list_of_values = [3, 6]
result = df.query("A in @list_of_values")
print(result)

Результат будет таким:

   A  B
1  6  2
2  3  3

В данном примере метод query позволяет отфильтровать DataFrame по значениям, представленным в списке list_of_values. Знак @ используется для доступа к переменным из внешнего пространства имен внутри строки запроса.

Вы можете использовать query, например:

b = df.query('a > 1 & a < 5')

Вы также можете использовать диапазоны таким образом:

b = df[(df['a'] > 1) & (df['a'] < 5)]

Этот код создает новый DataFrame b, который содержит строки из исходного DataFrame df, где значения в столбце 'a' находятся в диапазоне от 1 до 5 (не включая 1 и 5).

Извлечение данных с помощью pandas

Вот пример датафрейма:

    RPT_Date  STK_ID STK_Name  sales
0 1980-01-01       0   Arthur      0
1 1980-01-02       1    Beate      4
2 1980-01-03       2    Cecil      2
3 1980-01-04       3     Dana      8
4 1980-01-05       4     Eric      4
5 1980-01-06       5    Fidel      5
6 1980-01-07       6   George      4
7 1980-01-08       7     Hans      7
8 1980-01-09       8   Ingrid      7
9 1980-01-10       9    Jones      4

Существует несколько способов выборки или извлечения данных.

Использование .isin

Наиболее очевидный способ - использовать метод .isin(). Вы можете создать маску, которая даст серию значений True/False, которую можно применить к датафрейму:

mask = df['STK_ID'].isin([4, 2, 6])

mask
0    False
1    False
2     True
3    False
4     True
5    False
6     True
7    False
8    False
9    False
Name: STK_ID, dtype: bool

df[mask]
    RPT_Date  STK_ID STK_Name  sales
2 1980-01-03       2    Cecil      2
4 1980-01-05       4     Eric      4
6 1980-01-07       6   George      4

Использование масок - это простой способ решения задачи, хотя он не всегда эффективен с точки зрения скорости и использования памяти.

С помощью индексации

Установив индекс в столбец STK_ID, можно использовать встроенный метод среза pandas .loc.

df.set_index('STK_ID', inplace=True)
         RPT_Date STK_Name  sales
STK_ID                           
0      1980-01-01   Arthur      0
1      1980-01-02    Beate      4
2      1980-01-03    Cecil      2
3      1980-01-04     Dana      8
4      1980-01-05     Eric      4
5      1980-01-06    Fidel      5
6      1980-01-07   George      4
7      1980-01-08     Hans      7
8      1980-01-09   Ingrid      7
9      1980-01-10    Jones      4

df.loc[[4, 2, 6]]
         RPT_Date STK_Name  sales
STK_ID                           
4      1980-01-05     Eric      4
2      1980-01-03    Cecil      2
6      1980-01-07   George      4

Этот способ более быстрый, даже если установка индекса занимает некоторое время. Однако он экономит время, если вы хотите выполнять несколько таких запросов.

Объединение датафреймов

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

stkid_df = pd.DataFrame({"STK_ID": [4,2,6]})
df.merge(stkid_df, on='STK_ID')
   STK_ID   RPT_Date STK_Name  sales
0       2 1980-01-03    Cecil      2
1       4 1980-01-05     Eric      4
2       6 1980-01-07   George      4

Примечание

Все вышеуказанные методы работают, даже если есть несколько строк с одинаковым значением 'STK_ID'.