Создать новый столбец, используя для цикла Панды

Question

У меня есть следующий набор данных:

d = {'person':[1,1,1,1,1,1],'id':['-8','-5','-4','-3','-3','-2'],'obs': 
['A','B','C','D','E','F']}
df_start = pd.DataFrame(data=d)

Необходимо создать выходной набор данных, например:

d_end = {'id':[-8,-8,-5,-8,-5,-4,-5,-4,-3,-3,-5,-4,-3,-3],
'obs':['A','A','B','A','B','C','B','C','D','E','B','C','D','E'],
'id_group':[-8,-5,-5,-4,-4,-4,-3,-3,-3,-3,-2,-2,-2,-2]}
df_end = pd.DataFrame(data=d_end)

Я пытаюсь сгруппировать строки, используя новый столбец с именем id_group, который создается путем сравнения значений идентификаторов по строкам. Один идентификатор будет принадлежать его собственной группе идентификаторов. Идентификатор будет принадлежать другой группе идентификаторов, если (id + 4) больше или равен идентификатору в другой строке.

Не удалось продвинуться слишком далеко, используя цикл for, пытаясь это сделать, очень открыт для предложений

Answer 1

Чтобы сделать это без циклов, выполните одно большое слияние, а затем поднабор:

df_start['id'] = df_start['id'].astype(int)
df_end = df_start.assign(dummy=1).merge(df_start[['id']].assign(dummy=1), 
                                        on='dummy',
                                        suffixes=['','_gp']).drop(columns='dummy')

df_end = df_end[df_end.id.between(df_end.id_gp - 4, df_end.id_gp)]

Выход:

    person  id obs  id_gp
0        1  -8   A     -8
1        1  -8   A     -5
2        1  -8   A     -4
7        1  -5   B     -5
8        1  -5   B     -4
9        1  -5   B     -3
10       1  -5   B     -3
11       1  -5   B     -2
14       1  -4   C     -4
15       1  -4   C     -3
16       1  -4   C     -3
17       1  -4   C     -2
21       1  -3   D     -3
22       1  -3   D     -3
23       1  -3   D     -2
27       1  -3   E     -3
28       1  -3   E     -3
29       1  -3   E     -2
35       1  -2   F     -2

---

В цикле создайте подмножества для всех значений в пределах 4 id. Объединить их всех вместе.

df_start['id'] = df_start['id'].astype(int)

l = []
for id_gp in df_start.id.unique():
    l.append(df_start[df_start.id.between(id_gp-4, id_gp)].assign(id_gp = id_gp))

df_end = pd.concat(l, ignore_index=True)