使用Pandas识别同用户±10分钟窗口内创建的ID优化方案咨询

针对你处理10万条数据时嵌套循环耗时过长的问题，我们可以通过分组+二分查找的方式将时间复杂度从O(n²)降到O(n log n)，大幅提升运行效率。以下是具体实现思路和代码：

核心思路 #

1. 按用户分组：不同用户的记录无需互相比较，分组后将问题拆解为多个小任务，减少计算量。
2. 时间序列排序：对每个用户的记录按时间排序，便于后续用二分查找快速定位时间窗口范围。
3. 二分查找窗口边界：利用bisect模块快速找到每条记录±10分钟窗口内的首尾索引，计算窗口内的记录数。

代码实现 #

import pandas as pd
import bisect

# 1. 预处理数据：转换日期为datetime类型
data = pd.DataFrame({
    'ID': [1,2,3,4,5],
    'user': ['abc','abc','def','def','abc'],
    'Date': ['2020-02-02 10:01:00','2020-02-02 10:09:00','2020-02-02 10:17:00','2020-02-02 11:00:00','2020-02-02 10:17:00']
})
data['Date'] = pd.to_datetime(data['Date'])

# 2. 定义处理单个用户组的函数
def count_window_records(group):
    # 对组内记录按时间排序，保留原始索引以便后续合并
    sorted_group = group.sort_values('Date').reset_index()
    # 将日期转换为数值型时间戳（方便bisect比较）
    timestamps = sorted_group['Date'].astype('int64').values
    # 10分钟对应的纳秒数（datetime64的单位是纳秒）
    ten_min_ns = 10 * 60 * 10**9
    
    # 计算每条记录的时间窗口边界
    left_bounds = timestamps - ten_min_ns
    right_bounds = timestamps + ten_min_ns
    
    # 用二分查找找到每个边界在排序后时间序列中的位置
    left_indices = [bisect.bisect_left(timestamps, lb) for lb in left_bounds]
    right_indices = [bisect.bisect_right(timestamps, rb) for rb in right_bounds]
    
    # 计算窗口内的其他同用户记录数（排除自身）
    # 如果需要包含自身，将公式改为 r - l
    sorted_group['CNT'] = [r - l - 1 for l, r in zip(left_indices, right_indices)]
    
    # 恢复原始数据的顺序
    return sorted_group.set_index('index').sort_index()

# 3. 按用户分组应用函数，得到结果
result = data.groupby('user', group_keys=False).apply(count_window_records)

# 输出结果
print(result[['ID', 'user', 'Date', 'CNT']])

结果解释 #

运行上述代码后，输出结果与你的期望输出逻辑一致（注：若需要包含自身的计数，只需将r - l -1改为r - l即可）：

为什么这个方法更快？ #

• 二分查找的高效性：bisect的查找操作是O(log m)（m为单个用户的记录数），远快于嵌套循环的O(m²)。
• 分组减少计算量：仅在同用户内部进行比较，避免了跨用户的无效计算。
• 向量化预处理：将日期转换为数值型时间戳，避免了循环中频繁的datetime对象比较。

备选方案：使用Pandas Rolling时间窗口 #

如果偏好纯Pandas API，也可以用滚动时间窗口实现：

data_sorted = data.sort_values(['user', 'Date'])
# 按用户分组，使用20分钟的时间窗口（前后各10分钟），并设置center=True
data_sorted['CNT'] = data_sorted.groupby('user')['Date'].rolling(window='20T', center=True, closed='both').count().reset_index(level=0, drop=True) - 1
# 恢复原始顺序
result = data_sorted.sort_index()

这个方法同样高效，但需要注意窗口大小设置为20分钟（因为center=True时，窗口会以当前时间为中心前后各10分钟）。

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