присвоение значений первым трем строкам каждой группы

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

Если вы не дублировали индекс, вы можете создать идентификатор строки для каждой группы, отфильтровывать идентификатор, который больше, чем 3, а затем назначить его обратно к кадру данных:

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