выбор строк на pandas кадре данных в зависимости от условий

Question

У меня есть следующий код:

import pandas as pd
import random


a = [random.randint(0, 1) for i in range(30)]
b = [random.randint(0, 1) for i in range(30)]

print(a)
print(b)

df = pd.DataFrame([a, b])
df = df.T

columns = ['column1', 'column2']
df.columns = columns
print(df)

, который создает фрейм данных, хранящийся в переменной 'df'. Он состоит из 2 столбцов (column1 и column2), заполненных случайными 0 и 1.

Это вывод, полученный при запуске программы (если вы попытаетесь запустить ее, вы не получите точно такой же результат из-за генерации randomint).

    column1  column2
0         0        1
1         1        0
2         0        1
3         1        1
4         0        1
5         1        1
6         0        1
7         1        1
8         1        0
9         0        1
10        0        0
11        1        1
12        1        1
13        0        1
14        0        0
15        0        1
16        1        1
17        1        1
18        0        1
19        1        0
20        0        0
21        1        0
22        0        1
23        1        0
24        1        1
25        0        0
26        1        1
27        1        0
28        0        1
29        1        0

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

    column1  column2
2         0        1
3         1        1
4         0        1
5         1        1
6         0        1
7         1        1

11        1        1
12        1        1
13        0        1

15        0        1
16        1        1
17        1        1
18        0        1

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

I хотел бы сделать это следующим образом.

filter1 = (some boolean condition) &/| (maybe some other stuff)
final_df = df[filter1]

Спасибо

Answer 1

Мы можем использовать GroupBy.transform.

n = 3
blocks = df['column2'].ne(df['column2'].shift()).cumsum()
m1 = (df.groupby(blocks)['column2']
        .transform('size').ge(n))
m2 = df['column2'].eq(1)
df_filtered = df.loc[m1 & m2]
# Alternative without df['column2'].eq(1)
#df_filtered = df.loc[m1.mul(df['column2'])]
print(df_filtered)

Выход

    column1  column2
2         0        1
3         1        1
4         0        1
5         1        1
6         0        1
7         1        1

11        1        1
12        1        1
13        0        1

15        0        1
16        1        1
17        1        1
18        0        1

Если столбец 2 действительно содержит только 1 и 0 в Ваш оригинальный DataFrame, тогда мы можем использовать transform('sum') вместо transform('size')

---

, у блоков будет новое значение каждый раз, когда значение в column2 изменяется

print(blocks)
0      1
1      2
2      3
3      3
4      3
5      3
6      3
7      3
8      4
9      5
10     6
11     7
12     7
13     7
14     8
15     9
16     9
17     9
18     9
19    10
20    10
21    10
22    11
23    12
24    13
25    14
26    15
27    16
28    17
29    18
Name: column2, dtype: int64

Alternative

Я часто использую этот код в своих проектах, и я пришел к выводу, что, как правило, он может быть немного быстрее Series.map + Series.value_counts . Разница в производительности между этими двумя методами никогда не будет большой, и вы можете выбрать тот, который вы хотите. Но я обычно использую последний, который я объяснил, и я думаю, что стоит упомянуть об этом

%%timeit
m1 = blocks.map(blocks.value_counts().ge(n))
1.41 ms ± 122 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)


%%timeit
m1 = (df.groupby(blocks)['column2']
        .transform('size').ge(n))
2.12 ms ± 226 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)