Pandas：使用太多内存和条件滚动计数

Question

在给定固定窗口长度的情况下，我正在尝试对另一列中指定的组中出现的观察结果进行滚动计数。最好用一个例子来解释：

df = pd.DataFrame({'B': ['X', 'X' , 'Y', 'X', 'Y', 'Y', 'X', 'X', 'Y', 'Y', 'X', 'Y'],
                   'group': ["IT", "IT", "IT", "MV", "MV", "MV", "IT", "MV", "MV", "IT", "IT", "MV"]})

for i in df['B'].unique():
    df.loc[df['B']==i, 'count'] = df.where(df['B'].eq(i)).groupby(df['group'])['B'].transform(lambda x: x.rolling(3, min_periods=1).count().shift(fill_value=0))
print(df)

    B group  count
0   X    IT    0.0
1   X    IT    1.0
2   Y    IT    0.0
3   X    MV    0.0
4   Y    MV    0.0
5   Y    MV    1.0
6   X    IT    2.0
7   X    MV    1.0
8   Y    MV    2.0
9   Y    IT    1.0
10  X    IT    1.0
11  Y    MV    2.0

如上所述，我们按“组”分组，并对 B 列中的“X”和“Y”进行滚动计数，窗口长度 = 3。如果“X”是当前行，那么我们计算次数“ X' 出现在组 'group' 中的前 3 个观察值中，不包括当前行的计数（因此向后移动 period=1）。

但是，在处理大型数据集时，此代码速度较慢且占用的内存过多。感谢在这方面的改进。

Answer 1

我为您的问题开发了另一种解决方案，它基于 group-by 和 one-hot 编码 (get_dummy) 的使用。

代码如下：

df = pd.DataFrame({'B': ['X', 'X' , 'Y', 'X', 'Y', 'Y', 'X', 'X', 'Y', 'Y', 'X', 'Y'],
                   'group': ["IT", "IT", "IT", "MV", "MV", "MV", "IT", "MV", "MV", "IT", "IT", "MV"]})

# add a one-hot encoding to the dataframe. 
t = pd.concat([df, pd.get_dummies(df.B)], axis=1)

t.index.name = "inx"

# do a rolling sum of 4. It's the past 3, plus 1. 
t = t.groupby("group").rolling(4, min_periods = 1).sum()
t = t.reset_index().set_index("inx").sort_index()

# remove the extra '1' from the rolling result. 
t.loc[:, ["X", "Y"]] = t.loc[:, ["X", "Y"]] - 1

# merge back the results with the original dataframe. 
t = pd.concat([df, t[["X", "Y"]]], axis=1)

# create a 'count' column which is based on the values of 'B'. 
t["count"] = t.lookup(t.index, t.B )

输出是：

     B group    X    Y  count
inx                          
0    X    IT  0.0 -1.0    0.0
1    X    IT  1.0 -1.0    1.0
2    Y    IT  1.0  0.0    0.0
3    X    MV  0.0 -1.0    0.0
4    Y    MV  0.0  0.0    0.0
5    Y    MV  0.0  1.0    1.0
6    X    IT  2.0  0.0    2.0
7    X    MV  1.0  1.0    1.0
8    Y    MV  0.0  2.0    2.0
9    Y    IT  1.0  1.0    1.0
10   X    IT  1.0  1.0    1.0
11   Y    MV  0.0  2.0    2.0

---

一条龙：

df['count'] = (pd.concat([df, df['B'].str.get_dummies()], axis=1)
                 .groupby('group')
                 .rolling(4, min_periods=1)
                 .sum()
                 .sort_index(level=1)
                 .reset_index(drop=True)
                 .lookup(df.index, df['B']) - 1)