Столбец прямого заполнения с ограничением на основе индекса

Question

Я хочу заполнить столбец вперед и указать лимит, но я хочу, чтобы лимит основывался на индексе, а не на простом числе строк, как позволяет лимит.

Например, скажем, у меня есть фрейм данных, заданный как:

df = pd.DataFrame({
    'data': [0.0, 1.0, np.nan, 3.0, np.nan, 5.0, np.nan, np.nan, np.nan, np.nan],
    'group': [0, 0, 0, 1, 1, 0, 0, 0, 1, 1]
})

, который выглядит как

In [27]: df
Out[27]:
   data  group
0   0.0      0
1   1.0      0
2   NaN      0
3   3.0      1
4   NaN      1
5   5.0      0
6   NaN      0
7   NaN      0
8   NaN      1
9   NaN      1

Если я группирую по столбцу group и выполняю прямую заливкув этой группе с limit=2, тогда мой результирующий фрейм данных будет

In [35]: df.groupby('group').ffill(limit=2)
Out[35]:
   group  data
0      0   0.0
1      0   1.0
2      0   1.0
3      1   3.0
4      1   3.0
5      0   5.0
6      0   5.0
7      0   5.0
8      1   3.0
9      1   NaN

Однако то, что я на самом деле хочу здесь, это только прямая заливка в строки, индексы которых находятся в пределах, скажем, 2 от первого индекса каждогогруппа, в отличие от следующих 2 строк каждой группы.Например, если мы просто посмотрим на группы в кадре данных:

In [36]: for i, group in df.groupby('group'):
    ...:     print(group)
    ...:
   data  group
0   0.0      0
1   1.0      0
2   NaN      0
5   5.0      0
6   NaN      0
7   NaN      0
   data  group
3   3.0      1
4   NaN      1
8   NaN      1
9   NaN      1

Я бы хотел, чтобы вторая группа здесь была заполнена только до индекса 4 --- не 8 и 9. NaN первой группывсе значения находятся в пределах 2 индексов от последних значений, отличных от NaN, поэтому они будут заполнены полностью.Результирующий кадр данных будет выглядеть так:

   group  data
0      0   0.0
1      0   1.0
2      0   1.0
3      1   3.0
4      1   3.0
5      0   5.0
6      0   5.0
7      0   5.0
8      1   NaN
9      1   NaN

FWIW в моем фактическом случае использования, мой индекс - DateTimeIndex (и он отсортирован).

В настоящее время у меня есть решение, которое работаеттребуя циклически проходить через фрейм данных, отфильтрованный по групповым индексам, создавать временной диапазон для каждого отдельного события со значением, отличным от NaN, на основе индекса, а затем объединять их.Но это слишком медленно, чтобы быть практичным.

Answer 1

IIUC

l=[]
for i, group in df.groupby('group'):
    idx=group.index
    l.append(group.reindex(df.index).ffill(limit=2).loc[idx])

pd.concat(l).sort_index()
   data  group
0   0.0    0.0
1   1.0    0.0
2   1.0    0.0
3   3.0    1.0
4   3.0    1.0
5   5.0    0.0
6   5.0    0.0
7   5.0    0.0
8   NaN    1.0
9   NaN    1.0

Данные тестирования

   data  group
0   0.0      0
1   1.0      0
2   1.0      1
3   3.0      0
4   NaN      1
5   22       0
6   NaN      1
7   5.0      0
8   NaN      1
9   NaN      1

Мой метод тестирования данных

   data  group
0   0.0    0.0
1   1.0    0.0
2   1.0    1.0
3   3.0    0.0
4   1.0    1.0
5  22.0    0.0
6   NaN    1.0# here not change , since the previous two do not have valid value for group 1 
7   5.0    0.0
8   NaN    1.0
9   NaN    1.0

Out out with unutbu

   data  group
0   0.0      0
1   1.0      0
2   1.0      1
3   3.0      0
4   1.0      1
5  22.0      0
6   1.0      1# miss match in here
7   5.0      0
8   NaN      1
9   NaN      1

Answer 2

import numpy as np
import pandas as pd
df = pd.DataFrame({
    'data': [0.0, 1.0, 1, 3.0, np.nan, 22, np.nan, 5, np.nan, np.nan],
    'group': [0, 0, 1, 0, 1, 0, 1, 0, 1, 1]})

df = df.reset_index()
df['stop_index'] = df['index'] + 2
df['stop_index'] = df['stop_index'].where(pd.notnull(df['data']))
df['stop_index'] = df.groupby('group')['stop_index'].ffill()
df['mask'] = df['index'] <= df['stop_index']
df.loc[df['mask'], 'data'] = df.groupby('group')['data'].ffill()
print(df)
#    index  data  group  stop_index   mask
# 0      0   0.0      0         2.0   True
# 1      1   1.0      0         3.0   True
# 2      2   1.0      1         4.0   True
# 3      3   3.0      0         5.0   True
# 4      4   1.0      1         4.0   True
# 5      5  22.0      0         7.0   True
# 6      6   NaN      1         4.0  False
# 7      7   5.0      0         9.0   True
# 8      8   NaN      1         4.0  False
# 9      9   NaN      1         4.0  False

# clean up df
df = df[['data', 'group']]
print(df)

выходы

   data  group
0   0.0      0
1   1.0      0
2   1.0      1
3   3.0      0
4   1.0      1
5  22.0      0
6   NaN      1
7   5.0      0
8   NaN      1
9   NaN      1

---

Это копирует индекс в столбец, затем создает второй stop_index столбец, который index увеличивается на размер окно (времени).

df = df.reset_index()
df['stop_index'] = df['index'] + 2

Затем он создает пустые строки в stop_index, чтобы соответствовать пустым строкам в data:

df['stop_index'] = df['stop_index'].where(pd.notnull(df['data']))

Затем он заполняет форвард stop_index для каждой группы:

df['stop_index'] = df.groupby('group')['stop_index'].ffill()

Теперь (наконец-то) мы можем определить желаемые mask - места, где мы на самом деле хотим заполнить вперед data:

df['mask'] = df['index'] <= df['stop_index']
df.loc[df['mask'], 'data'] = df.groupby('group')['data'].ffill()