Агрегирование окон в массивы в пандах DataFrame

Question

У меня есть слава данных, например:

df = pd.DataFrame({"a": [1,2,3], "b": [4,5,6], "c": [7,8,9]})

a | b | c
1 | 4 | 7
2 | 5 | 8
3 | 6 | 9

И я бы хотел получить такой:

a     | b     | c
[1,2] | [4,5] | [7,8]
[2,3] | [5,6] | [8,9]

Итак, я попробовал самое очевидное: df.rolling(2).apply(lambda values: np.array(values)), который, к сожалению, не работает, так как rolling().apply строго ожидает скаляр (float) в качестве возвращаемого типа.

Так что я играл с пониманием.

window = 2
df = pd.DataFrame({"a": [1,2,3], "b": [4,5,6], "c": [7,8,9]})
df = pd.DataFrame({column:[df[column].iloc[i-window:i].values for i in range(window, len(df)+1)] for column in df})

Это правильно, но выглядит уродливо и очень медленно. Также он теряет тип индекса, который раньше был датой (теперь int). Есть ли лучший, более чистый способ?

Answer 1

Одним из решений может быть использование zip между значениями столбцов кадра данных:

df.apply(lambda x: list(zip(x[:-1:], x[1::])),
         raw=True).apply(lambda x: list(map(list, x))
                         ).apply(pd.Series).T

        a       b       c
0  [1, 2]  [4, 5]  [7, 8]
1  [2, 3]  [5, 6]  [8, 9]