Быстрый цикл через фрейм данных Python с предыдущей ссылкой на строку

Question

Предположим, у меня есть датафрейм pandas с двумя столбцами: ID и Дни.DataFrame сортируется в порядке возрастания по обеим переменным.Например:

# Initial dataset
data = pd.DataFrame({'id': np.repeat([1, 2 ,3], 4),
                 'day': [1, 2, 10, 11, 3, 4, 12, 15, 1, 20, 21, 24]})

    id  day
0   1   1
1   1   2
2   1   10
3   1   11
4   2   3
5   2   4
6   2   12
7   2   15
8   3   1
9   3   20
10  3   21
11  3   24

Я хочу добавить третий столбец, в котором указывался бы номер "сеанса" для каждого дня ID *.Под «сессией» я подразумеваю последовательность дней с разницей менее 2 дней между днями одного сеанса.Например, последовательность 5,6,7 будет считаться одной сессией, а 5,6,9 будет рассматриваться как две сессии и должна быть помечена как 0, 0, 1, т. Е. Дни 5 и 6 были отнесены к сеансу № 0, а день 9 -на сессию № 1.Номера сессий должны начинаться с 0 для каждого нового идентификатора.

Другими словами, я хочу получить следующее:

    id  day session
0   1   1   0
1   1   2   0
2   1   10  1
3   1   11  1
4   2   3   0
5   2   4   0
6   2   12  1
7   2   15  2  
8   3   1   0
9   3   20  1
10  3   21  1
11  3   24  2

Для решения этой задачи я использую Basic для цикла.В этом цикле я итеративно выполняю все уникальные идентификаторы, затем подгруппирую блок данных из начального набора данных и назначаю номера сеансов для каждого дня определенного идентификатора.У меня проблема - поскольку исходный набор данных состоит из миллионов строк - цикл занимает много времени !Например, на 1 млн строк мой цикл тратит около минуты, что слишком много.

Как улучшить скорость? Любой метод хорош!Если вы знаете, как достичь желаемого результата, например, с помощью некоторых манипуляций с матрицей, которые сократят время - тоже хорошо ...

Мой код для цикла:

# Get sessions for every id
sessions = []
for i in data.id.unique():
    id_data = data['day'][data['id']==i].reset_index(drop=True)
    for ind in id_data.index:
        if ind == 0:
            temp = [0]
        elif ((id_data[ind] - id_data[ind - 1]) < 2):
            temp.append(temp[ind - 1])
        else:
            temp.append(temp[ind - 1] + 1)
    sessions.extend(temp)

# Add sessions to the table
data['session'] = sessions 

Answer 1

Вы можете суммировать логическое значение

data.groupby('id').day.apply(lambda x : x.diff().gt(1).cumsum())
Out[614]: 
0     0
1     0
2     1
3     1
4     0
5     0
6     1
7     2
8     0
9     1
10    1
11    2
Name: day, dtype: int32

Answer 2

Мы можем воспользоваться тем, что ваши данные отсортированы, чтобы исключить fillna, сократить два groupby вызова до одного и устранить необходимость в apply.

df['session'] = df.day.diff().ge(2)
df['session'] = df.groupby('id').session.cumsum()

df
    id  day  session
0    1    1      0.0
1    1    2      0.0
2    1   10      1.0
3    1   11      1.0
4    2    3      0.0
5    2    4      0.0
6    2   12      1.0
7    2   15      2.0
8    3    1      0.0
9    3   20      1.0
10   3   21      1.0
11   3   24      2.0

В свою очередь, "session" будет столбцом с плавающей точкой.

Answer 3

Вы можете использовать groupby() дважды с np.where(), diff() и cumsum():

data['session'] = np.where(data.groupby('id')['day'].diff().fillna(0)>1, 1, 0)
data['session'] = data.groupby('id')['session'].cumsum()

Выход:

    id  day  session
0    1    1        0
1    1    2        0
2    1   10        1
3    1   11        1
4    2    3        0
5    2    4        0
6    2   12        1
7    2   15        2
8    3    1        0
9    3   20        1
10   3   21        1
11   3   24        2