如何为DataFrame中连续重复的gene_id添加别名后缀

我来帮你搞定这个需求！你要的是给连续重复的gene_id添加递增的_TE1/_TE2这类后缀，唯一出现的基因保持原名对吧？用pandas的矢量化分组+累计计数就能轻松实现，比你写的循环高效多了，而且一步到位~

完整解决方案代码 #

import pandas as pd

# 你的原始输入数据
df_genes_data = {'gene_id': ['g0', 'g1', 'g1', 'g2', 'g3', 'g4', 'g4', 'g4']}
df_genes = pd.DataFrame.from_dict(df_genes_data)

# 1. 标记连续相同gene_id的分组：连续相同的会被分到同一个组
# ne()判断当前行和上一行是否不同，cumsum()将布尔值转成分组编号
group_key = df_genes['gene_id'].ne(df_genes['gene_id'].shift()).cumsum()

# 2. 每个分组内从1开始累计计数
df_genes['seq_num'] = df_genes.groupby(group_key).cumcount() + 1

# 3. 预计算每个分组的元素数量（判断是否需要加后缀）
group_sizes = df_genes.groupby(group_key)['gene_id'].transform('size')

# 4. 拼接后缀：仅当分组元素>1时添加，否则保留原基因名
df_genes['gene_id'] = df_genes.apply(
    lambda row: f"{row['gene_id']}_TE{row['seq_num']}" if group_sizes[row.name] > 1 else row['gene_id'],
    axis=1
)

# 清理临时列
df_genes = df_genes.drop(columns='seq_num')

# 输出结果
print(df_genes.to_string())

运行这段代码后，你会得到完全符合预期的输出：

gene_id
0      g0
1  g1_TE1
2  g1_TE2
3      g2
4      g3
5  g4_TE1
6  g4_TE2
7  g4_TE3

代码步骤解释 #

1. 创建连续分组键：df['gene_id'].ne(df['gene_id'].shift()).cumsum() 会给每一组连续相同的gene_id分配唯一编号，比如示例里的g1两行是同一组，g4三行是同一组，单个基因各自成组。
2. 组内累计计数：groupby(group_key).cumcount()+1 让每个组内的元素从1开始递增计数，这就是后缀里的数字来源。
3. 判断是否加后缀：group_sizes 存储了每个分组的元素数量，只有当分组大小>1（即连续重复）时，才拼接_TE{序号}后缀，否则保持原基因名不变。

对你现有尝试的补充 #

你之前用循环提取连续重复基因的思路是对的，但pandas的矢量化操作能帮你省去手动遍历的麻烦，而且处理大数据集时速度会快很多，还能直接完成后缀添加的全流程，不用再单独处理提取到的重复基因列表~

备注：内容来源于stack exchange，提问作者emor