如何按ID分组根据Binary列生成符合规则的wanted(index)索引列

问题场景 #

现有按ID维度存储的包含Binary二值变量的数据集，样例结构如下：

ID Binary wanted(index)
1 0 0
1 0 0
1 1 0
1 0 1
1 0 2
1 0 3
2 0 0
2 1 0
2 0 1
3 1 0
3 0 1
......................

已知仅存在ID、Binary两列，需要生成第三列目标字段wanted(index)，索引生成规则如下：

• 按ID独立分组处理，跨ID不共享计数
• 每个分组内，1值出现前的所有0对应的索引值统一为0
• 遍历到1值时，当前位置索引记为0，1值之后出现的每一个0依次递增1赋值索引
• 若分组内首行就是1值（如样例中ID=3的分组），直接从首行开始执行索引计数

实现方案 #

1. Python Pandas 实现（本地表格处理场景） #

核心逻辑为先标记每个分组内首次出现1的分界点，分界点前全赋值0，分界点后按行顺序累计计数即可，代码可直接复用：

import pandas as pd

# 替换为你的数据读取逻辑，比如pd.read_csv()
df = pd.DataFrame({
    'ID': [1,1,1,1,1,1,2,2,2,3,3],
    'Binary': [0,0,1,0,0,0,0,1,0,1,0]
})

# 标记每个ID分组内是否已经出现过1
df['has_seen_1'] = df.groupby('ID')['Binary'].cummax() == 1
# 对已经出现1的分组区间，按顺序生成递增序号
df.loc[df['has_seen_1'], 'wanted(index)'] = df[df['has_seen_1']].groupby('ID').cumcount()
# 未出现1的区间统一赋值0
df['wanted(index)'] = df['wanted(index)'].fillna(0).astype(int)
# 移除临时列
df = df.drop(columns=['has_seen_1'])

运行后输出结果和样例完全匹配：

ID  Binary  wanted(index)
0    1       0              0
1    1       0              0
2    1       1              0
3    1       0              1
4    1       0              2
5    1       0              3
6    2       0              0
7    2       1              0
8    2       0              1
9    3       1              0
10   3       0              1

2. SQL 实现（数据库数仓处理场景） #

支持窗口函数的数据库（MySQL 8.0+、Hive、Spark SQL、ClickHouse等）都可以用如下逻辑实现：

WITH tmp AS (
    SELECT
        ID,
        Binary,
        -- 标记当前行之前（含当前行）是否已经出现过1，注意必须指定正确的排序字段
        MAX(Binary) OVER(
            PARTITION BY ID 
            ORDER BY 你的排序字段 -- 替换为表中能确定行顺序的字段，比如自增id、采集时间
            ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW
        ) AS has_seen_1
    FROM 你的表名
)
SELECT
    ID,
    Binary,
    CASE
        WHEN has_seen_1 = 0 THEN 0
        ELSE ROW_NUMBER() OVER(PARTITION BY ID, has_seen_1 ORDER BY 你的排序字段) - 1
    END AS `wanted(index)`
FROM tmp

注意：SQL实现必须指定能唯一确定行顺序的字段（如自增主键、数据生成时间戳），否则窗口函数的行遍历顺序不可控，会生成错误的索引值。

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