Python 如何选择和索引数据帧每组中的最高值？

我有一个包含多列的数据框架，每一列的组合描述一个实验（例如，多个超级标签，对于每个超级标签，多个片段具有不同的时间步数）。我想为所有实验设置每一集的最后一个时间步，但我不知道怎么做。我尝试了三种不同的方法，都是使用

.loc

和1）使用

.max（）.index

，2）使用

.idxmax（）

和3）使用

.tail（1）.index

，但都失败了（前两种方法都有我无法理解的异常，最后一种是错误的

这是我最简单的例子：

import numpy as np
import pandas as pd

np.random.seed(4)
def gen(t):
    results = []
    for episode_id, episode in enumerate(range(np.random.randint(2, 4))):
        for i in range(np.random.randint(2, 6)):
            results.append(
                {
                    "episode": episode_id,
                    "timestep": i,
                    "t": t,
                }
            )
    return pd.DataFrame(results)

df = pd.concat([gen("a"), gen("b")])

base_groups = ["t", "episode"]

df["last_timestep"] = False

print("Expected:")
print(df.groupby(base_groups).timestep.max())

#df.loc[df.groupby(base_groups).timestep.max().index, "last_timestep"] = True
#df.loc[df.groupby(base_groups).timestep.idxmax(), "last_timestep"] = True
df.loc[df.groupby(base_groups).tail(1).index, "last_timestep"] = True

print("Is:")
print(df[df.last_timestep])

df.groupby（base_groups）.timestep.max（）的输出正是我所期望的，选择了正确的行：

Expected:
t  episode
a  0          3
   1          4
b  0          2
   1          1
   2          4

但在过滤数据帧时，我得到的是：

Is:
   episode  timestep  t  last_timestep
2        0         2  a           True
3        0         3  a           True
4        1         0  a           True
8        1         4  a           True
2        0         2  b           True
3        1         0  b           True
4        1         1  b           True
8        2         3  b           True
9        2         4  b           True

不应选择第0行、第2行、第5行和第7行。
用于重复最大聚合值，并按列
时间步进行比较：

df["last_timestep"] = df.groupby(base_groups)['timestep'].transform(max).eq(df['timestep'])
print (df)
   episode  timestep  t  last_timestep
0        0         0  a          False
1        0         1  a          False
2        0         2  a          False
3        0         3  a           True
4        1         0  a          False
5        1         1  a          False
6        1         2  a          False
7        1         3  a          False
8        1         4  a           True
0        0         0  b          False
1        0         1  b          False
2        0         2  b           True
3        1         0  b          False
4        1         1  b           True
5        2         0  b          False
6        2         1  b          False
7        2         2  b          False
8        2         3  b          False
9        2         4  b           True