基于Pandas优化用户登录时间间隔的DataFrame处理

问题背景 #

处理网站用户登录时间数据：数据存储在Pandas DataFrame中，第一列为时间间隔，其余列标记用户在该间隔内是否登录。需求是生成实际存在同时登录情况的用户组合列，规则为：

• 若某行中多个用户同时登录，仅保留规模最大的用户组合列标记为1，其所有子集组合列标记为0
• 跳过从未同时登录的用户组合（如B和C从未同时登录，则无需生成B,C或A,B,C这类包含该无效组合的列）

现有代码的核心问题是：先生成所有可能的用户组合（数量随用户列数呈指数级增长），再反向清理无效列，当用户列较多时（如20个用户对应百万级组合），计算和内存成本极高。

核心优化思路 #

• 只生成有效组合：从单用户列出发，逐步向上推导存在同时登录记录的更大组合，避免生成所有可能的组合
• 提前过滤无效候选：仅基于已有有效组合生成候选组合，跳过不可能存在同时登录的组合
• 向量运算替代循环：用Pandas内置的高效向量操作替代Python层面的循环，提升计算效率
• 按需标记子集：在生成有效组合后，从最大组合开始反向标记其子集列，减少重复遍历

优化后的代码实现 #

import pandas as pd
from itertools import combinations

def process_login_combinations(df):
    # 分离时间列与用户列
    time_col = 'fecha_hora'
    user_cols = [col for col in df.columns if col != time_col]
    df_result = df[[time_col]].copy()
    
    # 初始化有效组合：所有单用户列（至少有登录记录）
    valid_combinations = [(col,) for col in user_cols]
    df_result[user_cols] = df[user_cols]
    
    # 从2用户组合开始，逐步生成更大的有效组合
    max_possible_size = len(user_cols)
    current_size = 2
    while current_size <= max_possible_size and valid_combinations:
        # 生成候选组合，用set去重避免重复检查
        candidate_combs = set()
        # 基于已有current_size-1的组合，与单用户组合合并生成候选
        for base_comb in [c for c in valid_combinations if len(c) == current_size-1]:
            for single_user in [c for c in valid_combinations if len(c) == 1 and c[0] not in base_comb]:
                new_comb = tuple(sorted(base_comb + single_user))
                candidate_combs.add(new_comb)
        
        # 验证候选组合是否存在同时登录的行，仅保留有效组合
        for comb in candidate_combs:
            has_active = df[list(comb)].all(axis=1).any()
            if has_active:
                col_name = ','.join(comb)
                df_result[col_name] = df[list(comb)].all(axis=1).astype(int)
                valid_combinations.append(comb)
        
        current_size += 1
    
    # 从最大组合开始，将其对应子集列的活跃行标记为0
    valid_combinations_sorted = sorted(valid_combinations, key=lambda x: len(x), reverse=True)
    for comb in valid_combinations_sorted:
        comb_size = len(comb)
        if comb_size <= 1:
            continue
        col_name = ','.join(comb)
        active_rows = df_result[col_name] == 1
        if not active_rows.any():
            continue
        # 遍历所有子集并标记
        for r in range(1, comb_size):
            for sub_comb in combinations(comb, r):
                sub_col_name = ','.join(sub_comb)
                if sub_col_name in df_result.columns:
                    df_result.loc[active_rows, sub_col_name] = 0
    
    # 移除全0列
    df_result = df_result.loc[:, (df_result != 0).any(axis=0)]
    
    return df_result

代码优化点说明 #

1. 无效组合过滤：仅基于已有有效组合生成候选，跳过从未同时登录的组合，大幅减少需要处理的列数量
2. 重复计算规避：用set存储候选组合，避免重复检查同一组合
3. 性能提升：使用Pandas内置的all(axis=1)向量操作替代逐行循环，计算效率提升显著
4. 内存优化：仅在确认组合有有效登录记录时才添加列，减少不必要的内存占用

测试示例 #

# 生成示例数据
date_rng = pd.date_range(start='2024-05-13 15:52:00', end='2024-05-13 16:04:00', freq='min')
df = pd.DataFrame(date_rng, columns=['fecha_hora'])
df['A'] = 1
df['B'] = [1] * 6 + [0] * 5 + [1] * 2
df['C'] = [1] * 5 + [0] * 6 + [1] * 2
df['D'] = [1] * 4 + [0] * 7 + [1] * 2
df['E'] = [1] * 3 + [0] * 3 + [1] * 2 + [0] * 3 + [1]*2
df['F'] = [0] * 7 + [1] * 3 + [0] * 3
df['alfa'] = [0] * 10 + [1] * 1 + [0] * 2
df.loc[1, ['A', 'B', 'C']] = 1
df.loc[8, ['D', 'E', 'F']] = 1

# 使用优化后的函数处理
df_optimized = process_login_combinations(df)

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