按ID分组校验余额归零后恒为0规则的Pandas与SQL解决方案咨询

校验规则说明：用户任意月份余额归零后，后续所有月份余额不能为非0值；从未出现过余额为0的用户默认合规。最终输出合规/不合规用户数统计结果。

1. Python-Pandas 实现 #

实现逻辑：依赖分组累积标记法，无需逐用户遍历循环，适配万级用户量也能秒级处理

• 先将year-month字段转为可排序的日期格式，原生YYYY-MM格式字符串也可直接排序，转datetime类型更稳妥
• 按用户ID分组后，对每个用户的记录按时间升序排列
• 累积标记用户是否已经出现过余额为0的记录，一旦出现后标记永久为True
• 只要标记为True的记录中存在余额非0的情况，该用户即为不合规

代码示例：

import pandas as pd

# 导入原始数据集到df变量
# 处理年月字段转日期格式
df['dt'] = pd.to_datetime(df['year-month'], format='%Y-%m')
# 按用户ID、时间升序排序
df = df.sort_values(['user ID', 'dt']).reset_index(drop=True)
# 分组计算每个用户是否已出现过0余额的累积标记
# 浮点型避免精度误差可将x==0替换为abs(x) < 1e-6
df['has_zero'] = df.groupby('user ID')['balance'].transform(lambda x: (x == 0).cummax())
# 筛选不合规用户列表：出现0后又出现非0余额的用户
invalid_users = df.loc[(df['has_zero'] == True) & (df['balance'] != 0)]['user ID'].unique()
# 全量用户列表
all_users = df['user ID'].unique()
# 生成最终统计结果
result = pd.DataFrame([
    {'well_informed': 'YES', 'num_cases': len(set(all_users) - set(invalid_users))},
    {'well_informed': 'NO', 'num_cases': len(invalid_users)}
])
print(result)

2. SQL 实现 #

实现逻辑：用窗口函数完成累积标记计算，无需嵌套循环，兼容大部分主流数据库都支持该语法
假设表名为user_balance，代码示例：

with user_calc as (
    select 
        `user ID` as user_id,
        -- 累积标记到当前行是否出现过余额为0的记录
        -- 浮点型避免精度误差可将balance=0替换为abs(balance) < 1e-6
        max(case when balance = 0 then 1 else 0 end) over(partition by `user ID` order by `year-month` rows between unbounded preceding and current row) as zero_flag,
        balance
    from user_balance
),
invalid_users as (
    -- 筛选出归零后又出现非0余额的不合规用户
    select distinct user_id from user_calc where zero_flag = 1 and balance != 0
),
all_users as (
    select distinct `user ID` as user_id from user_balance
)
-- 合并统计结果
select 'YES' as well_informed, count(*) as num_cases from all_users where user_id not in (select user_id from invalid_users)
union all
select 'NO' as well_informed, count(*) as num_cases from invalid_users;

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