Python中实现带重置的累计计数并保留连续重复段最后项

问题描述 #

现有如下NumPy数组：

import numpy as np
v = np.array([
  0, 0, 0, 1, 3, 3, 1, 1, 1, 1, 1, 0, 2, 3, 2, 1, 1, 0, 0, 1, 3, 3,
  3, 2, 0, 0, 0, 0, 1, 1, 1, 2, 1, 0, 0, 0, 0, 0, 1, 2, 2, 1, 0, 0,
  1, 1, 1, 0, 0, 0, 1, 2, 2, 1, 0, 0, 1, 1, 1, 1, 1, 2, 1, 1, 2, 0,
  0, 1, 2, 2, 2, 2, 1, 1, 2, 2, 1, 0, 0, 0, 0, 0, 0, 0, 1, 1, 2, 2,
  2, 0, 0, 0, 0, 1])

已通过Pandas实现连续重复段长度计算：

import pandas as pd
df = pd.DataFrame(v, columns=['digit'])
df["seq_len"] = df.groupby(
    (df["digit"] != df["digit"].shift()).cumsum()
)["digit"].cumcount()+1

得到的中间结果示例：

digit  seq_len
0       0        1
1       0        2
2       0        3
3       1        1
4       3        1
..    ...      ...
89      0        1
90      0        2
91      0        3
92      0        4
93      1        1

现在需要实现保留每个连续重复段的最后一行（即该段seq_len最大的行），常规的duplicated或drop_duplicates无法满足分段去重需求，比如执行df.drop_duplicates(subset='digit', keep='last')会得到错误结果：

digit  seq_len
22      3        3
88      2        3
92      0        4
93      1        1

期望输出结果：

digit  seq_len
2       0        3
3       1        1
5       3        2
10      1        5
..    ...      ...
88      2        3
92      0        4
93      1        1

同时需要通过公式index - seq_len + 1计算每个连续段的真实起始索引，最终结果示例：

start_index  digit  seq_len
0               0        3
3               1        1
4               3        2
6               1        5
..            ...      ...
86              2        3
89              0        4
92              1        1

高效解决方案 #

方法1：分组取每组末尾行 #

基于已有的分组键，直接按分组取最后一行，同时计算起始索引：

def magic_function(df):
    # 生成连续重复段的分组键
    group_key = (df["digit"] != df["digit"].shift()).cumsum()
    # 按分组键取每个段的最后一行
    result = df.groupby(group_key).last()
    # 计算起始索引
    result["start_index"] = result.index - result["seq_len"] + 1
    # 可选：将start_index设为索引
    # result = result.set_index("start_index")
    return result

方法2：布尔掩码筛选（更高效） #

直接生成掩码标记每个连续段的末尾行，无需分组操作，适合大数据量：

def magic_function(df):
    # 生成掩码：当前行与下一行digit不同，或为最后一行
    mask = (df["digit"] != df["digit"].shift(-1)) | (df.index == len(df)-1)
    # 筛选目标行
    result = df[mask].copy()
    # 计算起始索引
    result["start_index"] = result.index - result["seq_len"] + 1
    # 可选：将start_index设为索引
    # result = result.set_index("start_index")
    return result

结果验证 #

调用函数后即可得到目标结果：

output = magic_function(df)
# 输出核心列
print(output[["digit", "seq_len"]])
# 输出带起始索引的结果
print(output.reset_index(drop=True).set_index("start_index"))

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