Замените приведенные ниже значения n строк указанного столбца c тем же значением

Question

У меня есть такой фрейм данных,

df
col1    col2
 1       D
 2       A
 3       H
 4       K
 5       G
 6       A
 7       K
 8       H
 9       B
10       S

Теперь, если значение col2 равно A, заполните следующие две строки (мы можем изменить число) с помощью A, поэтому результат будет выглядеть следующим образом:

df
col1    col2
 1       D
 2       A
 3       A
 4       A
 5       G
 6       A
 7       A
 8       A
 9       B
10       S

Я могу сделать это, используя для l oop и сравнивая одно с другим. Но время выполнения будет больше. Поэтому я ищу pandas ярлыки / pythoni c способ сделать это эффективно.

Answer 1

Используйте Series.where и Series.ffill с limit=2. Наконец, мы используем Series.fillna, чтобы завершить значения, не являющиеся A и не выходящие за пределы.

df['col2'] = df['col2'].where(df['col2'].eq('A')).ffill(limit=2).fillna(df['col2'])

Выход

   col1 col2 new_col2
0     1    D        D
1     2    A        A
2     3    H        A
3     4    K        A
4     5    G        G
5     6    A        A
6     7    K        A
7     8    H        A
8     9    B        B
9    10    S        S

Мы также можем использовать DataFrame.rolling

df.loc[df['col2'].eq('A').rolling(3, min_periods=0).max().astype(bool),'col2'] = 'A'

Answer 2

Вы можете попробовать использовать ffill с limit=2 и fillna назад

df['new_col2'] = df.col2.where(df.col2.eq('A')).ffill(limit=2).fillna(df.col2)

Out[164]:
   col1 col2 new_col2
0     1    D        D
1     2    A        A
2     3    H        A
3     4    K        A
4     5    G        G
5     6    A        A
6     7    K        A
7     8    H        A
8     9    B        B
9    10    S        S