Python Pandas: Как получить индексы строк, где значение в столбце соответствует заданному?

df.iloc[i] возвращает i-ю строку DataFrame df. Здесь i не относится к метке индекса, а представляет собой индекс, основанный на нуле.

В отличие от этого, атрибут index возвращает фактические метки индекса, а не числовые индексы строк:

df.index[df['BoolCol'] == True].tolist()

или, эквивалентно,

df.index[df['BoolCol']].tolist()

Разницу можно увидеть довольно ясно, экспериментируя с DataFrame, в котором используется не стандартный индекс, не равный числовой позиции строки:

df = pd.DataFrame({'BoolCol': [True, False, False, True, True]},
       index=[10,20,30,40,50])

In [53]: df
Out[53]: 
   BoolCol
10    True
20   False
30   False
40    True
50    True

In [54]: df.index[df['BoolCol']].tolist()
Out[54]: [10, 40, 50]

Если вы хотите использовать индекс,

In [56]: idx = df.index[df['BoolCol']]

In [57]: idx
Out[57]: Int64Index([10, 40, 50], dtype='int64')

то вы можете выбрать строки, используя loc, вместо iloc:

In [58]: df.loc[idx]
Out[58]: 
   BoolCol
10    True
40    True
50    True

[3 строки x 1 столбец]

Обратите внимание, что loc также может принимать булевы массивы:

In [55]: df.loc[df['BoolCol']]
Out[55]: 
   BoolCol
10    True
40    True
50    True

[3 строки x 1 столбец]

Если у вас есть булев массив mask, и вам нужны порядковые индексы, вы можете вычислить их с помощью np.flatnonzero:

In [110]: np.flatnonzero(df['BoolCol'])
Out[112]: array([0, 3, 4])

Используйте df.iloc для выбора строк по порядковому индексу:

In [113]: df.iloc[np.flatnonzero(df['BoolCol'])]
Out[113]: 
   BoolCol
10    True
40    True
50    True

Да, это можно сделать с помощью функции numpy.where(). Вот пример того, как можно использовать её для выбора строк в DataFrame в зависимости от значения булевой колонки:

import pandas as pd
import numpy as np

# Создаем DataFrame
df = pd.DataFrame({"gene_name": ['SLC45A1', 'NECAP2', 'CLIC4', 'ADC', 'AGBL4'], 
                   "BoolCol": [False, True, False, True, True]},
                  index=list("abcde"))

# Проверяем содержимое DataFrame
print(df)

# Используем np.where для получения индексов, где BoolCol == True
indices = np.where(df["BoolCol"] == True)

# Преобразуем индексы в список
select_indices = list(indices[0])

# Используем iloc для выбора строк по индексам
selected_rows = df.iloc[select_indices]

# Выводим отобранные строки
print(selected_rows)

Этот код создаст DataFrame и затем выберет строки, где колонка BoolCol равна True. Результатом будет следующий DataFrame:

  BoolCol gene_name
b    True    NECAP2
d    True       ADC
e    True     AGBL4

Если вам нужны индексы отобранных строк, вы можете получить их с помощью .index:

# Получаем индексы отобранных строк
index_list = df.iloc[select_indices].index
print(index_list.tolist())  # Это преобразует индекс в список

Вывод будет следующим:

['b', 'd', 'e']

Таким образом, метод numpy.where() очень удобен для фильтрации данных в DataFrame на основе булевых условий.

Если вам нужно использовать объект DataFrame только один раз, вы можете сделать это следующим образом:

df['BoolCol'].loc[lambda x: x==True].index

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

Простой способ — сбросить индекс DataFrame перед фильтрацией:

df_reset = df.reset_index()
df_reset[df_reset['BoolCol']].index.tolist()

Немного хакерский подход, но быстро!

Другой способ — использовать метод pipe(), чтобы получить индексы, где значения в столбце BoolCol равны True. С точки зрения производительности, этот метод так же эффективен, как и классическое индексирование с использованием [].

df['BoolCol'].pipe(lambda x: x.index[x])

Это особенно полезно, если BoolCol является результатом нескольких сравнений, и вы хотите использовать метод цепочки (method chaining), чтобы собрать все операции в одном конвейере.

Например, если вам нужно получить индексы строк, где значение в NumCol больше 0.5, значение в BoolCol равно True, и произведение значений NumCol и BoolCol больше 0, вы можете сделать это, оценив выражение через eval() и вызвав pipe() на результате, чтобы выполнить индексирование индексов.

df.eval("NumCol > 0.5 and BoolCol and NumCol * BoolCol > 0").pipe(lambda x: x.index[x])

¹: В следующем бенчмарке использовался DataFrame с 20 миллионами строк (в среднем отфильтровывалось половина строк) для получения их индексов. Метод цепочки с помощью pipe() показывает отличные результаты по сравнению с другими эффективными вариантами.

n = 20_000_000
df = pd.DataFrame({'NumCol': np.random.rand(n).astype('float16'), 
                   'BoolCol': np.random.default_rng().choice([True, False], size=n)})

%timeit df.index[df['BoolCol']]
# 181 ms ± 2.47 ms на цикл (среднее ± стандартное отклонение за 10 запусков, 1000 циклов каждый)

%timeit df['BoolCol'].pipe(lambda x: x.index[x])
# 181 ms ± 1.08 ms на цикл (среднее ± стандартное отклонение за 10 запусков, 1000 циклов каждый)

%timeit df['BoolCol'].loc[lambda x: x].index
# 297 ms ± 7.15 ms на цикл (среднее ± стандартное отклонение за 10 запусков, 1000 циклов каждый)

²: Для DataFrame с 20 миллионами строк, сконструированного так же, как в ¹ для бенчмарка, вы обнаружите, что предложенный здесь метод является самым быстрым вариантом. Он работает лучше, чем цепочка побитовых операторов, потому что по своей природе eval() выполняет несколько операций на большом DataFrame быстрее, чем векторизированные операции Python, и более эффективно использует память, чем query(), потому что, в отличие от query(), eval().pipe(...) не требует создания копии вырезанного DataFrame для получения его индекса.