为每个组的前三行分配值

Question

我正在尝试在pandas中编写以下逻辑代码，对于每组的前三行，我想创建一个应该具有值1（第1行），2（第2行），3（第3行）的变量。我在下面这样做，在下面的代码中，我没有创建一个新变量，因为我不知道如何做到这一点，所以我正在替换已经存在于数据集中的变量。虽然我的代码没有抛出错误，但它给了我非常奇怪的结果。

def func (i):
    data.loc[data.groupby('ID').nth(i).index,'date'] = i
func(1)

有什么建议吗？ 在此先感谢。

Answer 1

如果您没有重复索引，则可以为每个组创建一个行ID，过滤掉大于3的id，然后将其分配回数据框：

data['date'] = (data.groupby('ID').cumcount() + 1)[lambda x: x <= 3]

这会为每个ID 1,2,3提供前三行，超过3的行将包含NaN个值。

data = pd.DataFrame({"ID":[1,1,1,1,2,2,3,3,3]})
data['date'] = (data.groupby('ID').cumcount() + 1)[lambda x: x <= 3]
data