如何操作Python DataFrame实现用户订阅的索引、分组与分段？

嗨，咱们一步步来搞定这个需求！你用groupby的思路完全没问题，接下来咱们把这个思路落地，实现你要的筛选和分组功能～

第一步：先给每个用户标记订阅数量 #

首先咱们给原DataFrame加一个临时列，用来记录每个用户对应的订阅id总数。用transform方法能让每一行都带上所属用户的统计结果，不用拆分数据就能直接后续筛选：

# 给每行数据添加所属用户的订阅数
data1['sub_count'] = data1.groupby('user_id')['id'].transform('nunique')

这里nunique用来统计每个用户对应的唯一订阅id数量，transform保证原DataFrame的每一行都能拿到自己所属用户的统计值，而不是只返回分组后的汇总结果。

第二步：筛选单订阅/多订阅用户数据 #

有了sub_count列，筛选就变得非常简单：

提取仅关联一个订阅id的用户数据 #

single_sub_df = data1[data1['sub_count'] == 1]

提取关联多个订阅id的用户数据 #

multiple_sub_df = data1[data1['sub_count'] > 1]

如果你只需要去重的用户id列表（而非完整数据），可以这样处理：

# 单订阅用户id集合
single_sub_user_ids = single_sub_df['user_id'].unique()
# 多订阅用户id集合
multiple_sub_user_ids = multiple_sub_df['user_id'].unique()

封装成可复用的函数 #

如果需要多次使用这个逻辑，咱们可以把它封装成函数，还支持自定义列名适配不同的数据集：

def split_users_by_subscription(df, user_col='user_id', sub_col='id'):
    """
    将DataFrame按用户的订阅数量拆分，返回单订阅和多订阅的用户数据
    参数:
        df: 原始DataFrame
        user_col: 用户id列名，默认'user_id'
        sub_col: 订阅id列名，默认'id'
    返回:
        single_sub_df: 仅含单订阅用户的DataFrame
        multiple_sub_df: 仅含多订阅用户的DataFrame
    """
    # 复制原数据避免修改原始文件
    df_temp = df.copy()
    # 添加订阅数统计列
    df_temp['sub_count'] = df_temp.groupby(user_col)[sub_col].transform('nunique')
    
    # 筛选并移除临时列
    single_sub_df = df_temp[df_temp['sub_count'] == 1].drop(columns='sub_count')
    multiple_sub_df = df_temp[df_temp['sub_count'] > 1].drop(columns='sub_count')
    
    return single_sub_df, multiple_sub_df

使用的时候直接调用就行：

# 默认列名的情况
single_sub, multiple_sub = split_users_by_subscription(data1)

# 如果你的列名不是默认值，比如用户列叫'uid'，订阅列叫'sub_id'
# single_sub, multiple_sub = split_users_by_subscription(data1, user_col='uid', sub_col='sub_id')

关于你之前的groupby('user_id').groups #

你之前用这个方法得到的每个user_id对应的id索引集合是对的，但直接操作groups对象不如transform方便——后者能直接把统计结果映射回原DataFrame，让筛选和后续分析更顺畅。如果想从groups里提取信息，也可以这么做：

# 先获取每个用户的订阅数量字典
user_sub_counts = {user: len(ids) for user, ids in data1.groupby('user_id').groups.items()}
# 筛选单/多订阅用户id
single_sub_ids = [user for user, count in user_sub_counts.items() if count == 1]
multiple_sub_ids = [user for user, count in user_sub_counts.items() if count > 1]
# 再用id去原DataFrame筛选数据
single_sub_df = data1[data1['user_id'].isin(single_sub_ids)]

不过这种方法需要多一步字典推导，还是前面用transform的方式更简洁哦～

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