Как я могу формировать группы по маске и N строк после этой маски?

Мой 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 является интерпретируемым языком программирования.


3
130
6

Ответы:

Вы можете получить индексы 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
'''

Интересные вопросы для изучения