Pandas：统计DataFrame中0和1的模式数量并添加ID列

问题背景 #

给定如下Pandas DataFrame：

import pandas as pd

data = {'A': ['XYZ', 'XYZ', 'XYZ', 'XYZ', 'PQR', 'PQR', 'PQR', 'PQR', 'CVB', 'CVB', 'CVB', 'CVB'], 
        'B': ['2022-02-16 14:00:31', '2022-02-16 16:11:26', '2022-02-16 17:31:26',
              '2022-02-16 22:47:46', '2022-02-17 07:11:11', '2022-02-17 10:43:36', 
              '2022-02-17 15:05:11', '2022-02-18 18:06:12', '2022-02-19 09:05:46', 
              '2022-02-19 13:02:16', '2022-02-19 18:05:26', '2022-02-19 22:05:26'], 
        'C': [1,0,0,0,1,0,1,0,0,0,0,1]}
df = pd.DataFrame(data)
df['B'] = pd.to_datetime(df['B'])

需要实现：

• 统计列C中以1开头的连续模式的长度（即从当前1开始，到下一个1之前的所有行数，包含当前1），将统计值存入Count列，仅在C=1的行显示，其余行留空为NaN
• 为每个C=1的行生成唯一ID，格式为ABC_序号，仅在C=1的行显示，其余行留空

预期输出示例：

用户尝试的代码无法得到预期结果：

one_index = df[df['C'] == 1].index
zero_index = df[df['C'] == 0].index

df.loc[0, 'Count'] = len(df)
df.loc[one_index, 'ID'] = "ABC_1"

解决方案 #

可以通过标记连续分组、计算每组长度、生成ID三个步骤实现：

完整代码 #

import pandas as pd

data = {'A': ['XYZ', 'XYZ', 'XYZ', 'XYZ', 'PQR', 'PQR', 'PQR', 'PQR', 'CVB', 'CVB', 'CVB', 'CVB'], 
        'B': ['2022-02-16 14:00:31', '2022-02-16 16:11:26', '2022-02-16 17:31:26',
              '2022-02-16 22:47:46', '2022-02-17 07:11:11', '2022-02-17 10:43:36', 
              '2022-02-17 15:05:11', '2022-02-18 18:06:12', '2022-02-19 09:05:46', 
              '2022-02-19 13:02:16', '2022-02-19 18:05:26', '2022-02-19 22:05:26'], 
        'C': [1,0,0,0,1,0,1,0,0,0,0,1]}
df = pd.DataFrame(data)
df['B'] = pd.to_datetime(df['B'])

# 1. 标记以C=1开头的连续分组
# 当C=1时，标记为新分组的起点，累加分组ID
df['group_id'] = (df['C'] == 1).cumsum()

# 2. 计算每个分组的长度，仅赋值给C=1的行
df['Count'] = df.groupby('group_id')['group_id'].transform('count')
df.loc[df['C'] != 1, 'Count'] = pd.NA

# 3. 生成唯一ID，仅赋值给C=1的行
one_rows = df['C'] == 1
df.loc[one_rows, 'ID'] = 'ABC_' + (df[one_rows]['group_id'].astype(str))

# 移除临时的group_id列（可选）
df = df.drop('group_id', axis=1)

# 查看结果
print(df)

步骤解释 #

1. 标记分组：(df['C'] == 1).cumsum()会在每次遇到C=1时累加计数，这样每个以1开头的连续段会被分配同一个group_id。
2. 计算Count：通过groupby('group_id').transform('count')得到每个分组的行数，然后将非C=1的行的Count设为NaN。
3. 生成ID：筛选出C=1的行，用group_id拼接成ABC_序号的格式。

执行后即可得到符合预期的结果，其中CVB组的最后一行C=1的Count为1，ID为ABC_4，符合逻辑。

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