如何为Pandas DataFrame生成indexy列：值变更前取首索引+1

问题说明 #

原始DataFrame：

fruit
0    apple
1    apple
2    apple
3    banana
4    apple
5    pear

需要新增indexy列，规则为：连续相同fruit值的分组，indexy取值为该分组首个索引位置+1，直至fruit值变更，最终结果如下：

fruit    indexy
0    apple    1
1    apple    1
2    apple    1
3    banana   4
4    apple    5
5    pear     6

解决方案（Pandas实现） #

核心思路是先识别连续相同fruit的分组，再为每个分组提取首索引+1并广播到整组：

代码实现 #

import pandas as pd

# 构造原始数据
df = pd.DataFrame({'fruit': ['apple', 'apple', 'apple', 'banana', 'apple', 'pear']})

# 生成连续分组标识：通过比较当前行与上一行fruit是否不同，再累加得到分组ID
group_ids = df['fruit'].ne(df['fruit'].shift()).cumsum()

# 对每个分组，取首个索引+1，用transform广播到分组内所有行
df['indexy'] = df.groupby(group_ids)['fruit'].transform(lambda x: x.index[0] + 1)

# 查看结果
print(df)

代码解释 #

1. 生成连续分组标识：df['fruit'].ne(df['fruit'].shift())会生成布尔序列，当前行与上一行fruit不同时为True，否则为False；cumsum()将布尔值转为1/0累加，得到连续相同fruit的分组编号（比如前3个apple为组1，banana为组2，第4个apple为组3，pear为组4）。
2. 广播分组首索引+1：groupby(group_ids)按连续分组聚合，transform(lambda x: x.index[0]+1)提取每个分组的第一个索引值并加1，再将该值赋值给分组内的所有行，最终得到目标indexy列。

简化写法（无需中间变量） #

df['indexy'] = df.groupby(df['fruit'].ne(df['fruit'].shift()).cumsum())['fruit'].transform(lambda x: x.index[0] + 1)

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