python - 使用前一行引用快速循环 Python 数据框

Question

假设我有一个包含两列的 pandas 数据框:ID 和 Days。 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*day 提供一个“ session ”编号。 “ session ”是指一系列天数，一个 session 的天数之间的差异小于 2 天。例如，序列 5,6,7 将被视为一个 session ，而 5,6,9 将被视为两个 session 并应标记为 0, 0, 1，即第 5 天和第 6 天被引用到 session #0，而第 9 天被引用到 session #1。每个新 ID 的 session 编号应该从 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

为了解决这个任务，我使用了基本的 for 循环。在这个循环中，我迭代地遍历所有唯一 ID，然后从初始数据集中对数据 block 进行子集化，并为特定 ID 的每一天分配 session 编号。我遇到的问题 - 因为初始数据集是数百万行 - loop 需要很多时间!例如，对于 100 万行，我的循环花费了大约一分钟，这太多了。

如何提高速度？什么方法都好!如果你知道如何达到预期的结果，例如，通过一些 numpy 矩阵操作来减少时间 - 也很好......

我的循环代码:

# 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

你可以对 bool 求和

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