如何在DataFrame中按行方向计算窗口大小为n的滚动唯一值计数？

问题描述 #

给定DataFrame df，需要在行方向上，以指定窗口大小n为边界，计算每个滚动窗口内所有元素的唯一值数量，并将结果作为新列rolling_nunique添加到原DataFrame中。

输入数据 #

输入代码 #

import pandas as pd
import numpy as np

data = {'col_0':[7, 8, 9, 10, 11, 12],
        'col_1':[4, 5, 6, 7, 8, 9],
        'col_2':[2, 5, 8, 11, 14, 15],
        'col_3':[2, 6, 10, 14, 18, 21],
        'col_4':[7, 5, 7, 5, 7, 5],
        'col_5':[2, 6, 10, 14, 18, 21]}
df = pd.DataFrame(data)
print(df)

输入数据输出 #

col_0  col_1  col_2  col_3  col_4  col_5
0      7      4      2      2      7      2
1      8      5      5      6      5      6
2      9      6      8     10      7     10
3     10      7     11     14      5     14
4     11      8     14     18      7     18
5     12      9     15     21      5     21

预期输出（窗口大小n=2） #

输出结果 #

col_0  col_1  col_2  col_3  col_4  col_5  rolling_nunique
0      7      4      2      2      7      2                3
1      8      5      5      6      5      6                6
2      9      6      8     10      7     10                6
3     10      7     11     14      5     14                8
4     11      8     14     18      7     18                7
5     12      9     15     21      5     21               10

窗口逻辑说明 #

• 窗口0：仅包含第0行元素 [[7,4,2,2,7,2]]，唯一值为[2,4,7]，数量为3。
• 窗口1：包含第0行和第1行元素

  [[7,4,2,2,7,2]
   [8,5,5,6,5,6]]
  

  唯一值为[2,4,5,6,7,8]，数量为6。
• 窗口2：包含第1行和第2行元素

  [[8,5,5,6,5,6]
   [9,6,8,10,7,10]]
  

  唯一值为[5,6,7,8,9,10]，数量为6。
• 后续窗口以此类推，每个窗口包含当前行及前n-1行的所有元素。

解决方案 #

方法1：使用rolling结合apply #

利用Pandas的rolling方法按行创建窗口，再通过apply将窗口内的所有元素展平后计算唯一值数量：

def count_unique(window):
    # 将窗口内的多行数据展平为一维数组，计算唯一值数量
    return len(np.unique(window.to_numpy().flatten()))

# 设置窗口大小n
n = 2
df['rolling_nunique'] = df.rolling(window=n, min_periods=1).apply(count_unique, raw=False)
# 转换为整数类型（apply返回浮点数）
df['rolling_nunique'] = df['rolling_nunique'].astype(int)

print(df)

方法2：更高效的向量化实现（适用于大数据集） #

对于较大的DataFrame，apply效率较低，可以使用向量化方式处理：

n = 2
# 获取每个窗口的行索引范围
window_indices = [range(max(0, i-n+1), i+1) for i in range(len(df))]
# 遍历每个窗口，提取元素并计算唯一值数量
df['rolling_nunique'] = [len(np.unique(df.loc[idx].to_numpy().flatten())) for idx in window_indices]

print(df)

这两种方法都能得到符合预期的结果，方法2在数据量较大时性能更优。

内容的提问来源于stack exchange，提问作者Baron Legendre