Панды векторизовали способ создания групп размером n?

Question

Скажем, у меня огромный тензор формы (samples, timesteps, features), но я хочу развернуть его для выполнения groupby операций с пандами, как я могу соответствующим образом пометить все элементы размера n: n + в векторизованном виде? Медленное решение:

import pandas as pd
import numpy as np

df = pd.DataFrame(np.random.normal(0, 1, 500))
df["sample"] = np.nan

n_timesteps = 50
n_samples = len(df) // n_timesteps

size = n_timesteps
for i in range(n_samples):
    id0 = i * n_timesteps
    id1 = i * n_timesteps + n_timesteps
    df.loc[id0:id1, "sample"] = i

Answer 1

Назначить новый столбец с помощью целочисленного деления с помощью index:

#default RangeIndex
df['sample'] = df.index // n_timesteps

или с помощью массива 1d numpy, созданного arange:

df['sample'] = np.arange(len(df)) // n_timesteps