Мой DataFrame:
import pandas as pd
df = pd.DataFrame(
{
'a': [False, True, False, True, False, True, False, True, True, False, False],
}
)
Ожидаемый результат формирует такие группы:
a
1 True
2 False
3 True
a
5 True
6 False
7 True
a
8 True
9 False
10 False
Логика такова:
По сути, я хочу сформировать группы, где df.a == True
и две строки после этого. Например, чтобы создать первую группу, необходимо найти первую True
— строку 1
. Тогда первая группа — это строки 1, 2 и 3. Для второй группы необходимо найти следующий True
, которого нет в первой группе. Эта строка — строка 5
. Итак, вторая группа состоит из строк 5, 6 и 7. Это изображение проясняет суть:
И это моя попытка, которая не сработала:
N = 2
mask = ((df.a.eq(True))
.cummax().cumsum()
.between(1, N+1)
)
out = df[mask]
🤔 А знаете ли вы, что...
Python является интерпретируемым языком программирования.
Вы можете получить индексы True
с расстоянием 3:
import pandas as pd
df = pd.DataFrame(
{
'a': [False, True, False, True, False, True, False, True, True, False, False],
}
)
indices = df.index[df['a'] == True].tolist()
groups = [indices[0]]
for i in range(1, len(indices)):
if indices[i] - groups[-1] >= 3:
groups.append(indices[i])
for idx in groups:
end = min(idx + 3, len(df))
print(df.iloc[idx:end])
a
1 True
2 False
3 True
a
5 True
6 False
7 True
a
8 True
9 False
10 False
Я определил функцию mask_dataframe
, которая принимает DataFrame
и возвращает группы длины mask_len
+ 1 в новом DataFrame
, ища первый экземпляр, где из True
, и извлекая следующие mask_len
элементы.
Значение last
используется для пропуска элементов True
, которые уже были агрегированы в новый df.
import pandas as pd
df = pd.DataFrame(
{
'a': [False, True, False, True, False, True, False, True, True, False, False],
}
)
def mask_dataframe(df: pd.DataFrame, mask_len: int) -> pd.DataFrame:
last = 0 # keep track of which indices have been aggregated
new_df = pd.DataFrame() # create an empty df for the return data
for index, ele in enumerate(df['a']): # iterate over the input df
if ele and index > last: # if the element is True and hasn't been aggregated
last = index + mask_len # increment the index tracking
new_df = new_df._append(df.loc[index:last]) # and collect the data
return new_df
print(mask_dataframe(df, 2))
Это возвращает:
a
1 True
2 False
3 True
5 True
6 False
7 True
8 True
9 False
10 False
Другое возможное решение:
from itertools import accumulate
N = 3
# gets the indexes where each group starts
idx = np.unique(np.array(list(
accumulate(df.a[df.a].index, lambda x, y: y if (y - x) >= N else x)
)))
# creates the index for all groups
[df.iloc[i: i + N] for i in idx]
Выход:
[ a
1 True
2 False
3 True,
a
5 True
6 False
7 True,
a
8 True
9 False
10 False]
Чтобы создать группы, отфильтруйте фрейм данных, затем отфильтруйте еще раз, чтобы пропустить каждую вторую секунду True и переиндексировать обратно.
g = (df.loc[df['a'], 'a']
.iloc[::2]
.reindex(df.index, fill_value=False)
.cumsum()
)
groups = [val[:3] for name, val in df.groupby(g) if name != 0]
Конечный результат:
[ a
1 True
2 False
3 True,
a
5 True
6 False
7 True,
a
8 True
9 False
10 False]
Поскольку ваша проблема по своей сути итеративна, вы должны выполнить цикл.
Самый простой вариант IMO — использовать простой цикл Python:
def split(df, N=2):
i = 0
a = df['a'].to_numpy()
while i < len(df):
if a[i]:
yield df.iloc[i:i+N+1]
i+=N
i+=1
out = list(split(df))
Выход:
[ a
1 True
2 False
3 True,
a
5 True
6 False
7 True,
a
8 True
9 False
10 False]
Если вам нужна простая маска и уникальный DataFrame на выходе, вы можете повысить скорость с помощью numba:
from numba import jit
@jit(nopython=True)
def get_indices(a, N=2):
i = 0
out = []
while i < len(a):
if a[i]:
out.extend([True]*(N+1))
i += N+1
else:
out.append(False)
i += 1
return out[:len(a)]
out = df.loc[get_indices(df['a'].to_numpy())]
out['group'] = np.arange(len(out))//3
Выход:
a group
1 True 0
2 False 0
3 True 0
5 True 1
6 False 1
7 True 1
8 True 2
9 False 2
10 False 2
N = 2
относительные сроки:
N = 100
Примечание. Решение @Triky здесь не дает правильных результатов
относительные сроки
import pandas as pd
import numpy as np
# Example usage:
data = {
'a': [False, True, False, True, False, True, False, True, True, False, False],
}
df = pd.DataFrame(data)
print(df)
'''
a
0 False
1 True
2 False
3 True
4 False
5 True
6 False
7 True
8 True
9 False
10 False
'''
# Find indices where df['a'] is True
true_indices = np.flatnonzero(df['a'].values)
# Initialize the group column with NaNs
gr_column = np.full(len(df),np.nan)
current_gr = 0
for idx in true_indices :
if np.isnan(gr_column[idx]):
end_idx = min(idx + 3 ,len(df))
gr_column[idx:end_idx] = current_gr
current_gr += 1
print(gr_column)#[nan 0. 0. 0. nan 1. 1. 1. 2. 2. 2.]
df['group'] = gr_column
df = df.dropna().reset_index(drop=True)
print(df)
'''
a group
0 True 0.0
1 False 0.0
2 True 0.0
3 True 1.0
4 False 1.0
5 True 1.0
6 True 2.0
7 False 2.0
8 False 2.0
'''