pandas dataframe, группа, скользящее окно, затем статистика

Question

У меня есть DataFrame с двумя столбцами, A и B, все целые числа. В А они повторяются. Я группирую их, затем сортирую по ним, затем хочу применить скользящее окно над A, чтобы сгруппировать элементы в столбце B. То, как они сгруппированы, не имеет решающего значения, что-нибудь подойдет, тогда мне придется сделать над ними ряд вещей, поэтому мне, скорее всего, нужно будет преобразовать эту коллекцию во что-то другое (pandas.Series, вероятно, будут лучшими объединяющими группами. в каждом окне).

Я думаю, мне нужно начать с чего-то вроде этого:

df.groupby('A').rolling(w)

а как тогда получить элементы?

Я пытался использовать apply, но я получаю только одно окно в виде массива, а не окно над группами.

Пример:

In [1]: import pandas as pd
   ...: import numpy as np
   ...: import random
   ...: 
   ...: random.seed(2)
   ...: 
   ...: indexes = [i for i in range(1,100) for _ in range(10)]
   ...: dfi = pd.DataFrame({'A': indexes, 'B': [random.randint(1,99) for e in indexes]})
   ...: print(dfi.head()) # this is an input example
   A   B
0  1   8
1  1  12
2  1  11
3  1  47
4  1  22
In [2]: result = []
   ...: w = 3
   ...: for i in range(1,100):
   ...:     result.append({'A': i, 'B': np.array([e for j, e in dfi.values if abs(i-j) < w or abs(min(i,j)+99-max(i,j)
   ...: ) < w])})
   ...: dfo = pd.DataFrame(result) # this is the expected output, to be obtained with groupby and rolling operations
   ...: print(dfo.head())
   A                                                  B
0  1  [8, 12, 11, 47, 22, 95, 86, 40, 33, 78, 28, 78...
1  2  [8, 12, 11, 47, 22, 95, 86, 40, 33, 78, 28, 78...
2  3  [8, 12, 11, 47, 22, 95, 86, 40, 33, 78, 28, 78...
3  4  [28, 78, 5, 75, 88, 21, 56, 82, 51, 93, 66, 48...
4  5  [66, 48, 70, 57, 65, 35, 5, 4, 47, 60, 41, 49,...

Обратите внимание, что в каждом массиве 50 элементов, по 10 для каждой группы, а затем размер окна 5 (3 радиуса, не включительно)

Answer 1

Вот решение с использованием агрегата:

dfo = dfi.groupby('A').agg(lambda x: list(x))
dfo.head()
                                               B
    A                                           
    1    [8, 12, 11, 47, 22, 95, 86, 40, 33, 78]
    2    [28, 78, 5, 75, 88, 21, 56, 82, 51, 93]
    3     [66, 48, 70, 57, 65, 35, 5, 4, 47, 60]
    4    [41, 49, 55, 68, 22, 72, 23, 31, 30, 4]
    5   [23, 42, 23, 18, 66, 66, 47, 66, 87, 72]

И вы можете вернуть Series вместо DataFrame, обратившись к столбцу «B» в объекте DataFrameGroupBy, например:

dfo = dfi.groupby('A')['B'].apply(lambda x: list(x))