活用 shift，10 秒搞定 500 万

作者: 呆鸟的简书 | 来源:发表于2020-06-12 00:07 被阅读0次

Pandas 百问百答第 011 篇。

呆鸟云：“因受疫情影响，断更了好久，不过大家不用担心，只要呆鸟有心得，就会继续与大家分享，这个号不会黄的。这次先来一篇优化 pandas 计算效率的文章，希望大家喜欢。”

话说，前几天凤哥找我说要对几百万条数据先分组，再根据条件为每组里的数据加上序号, 据他说用循环的速度太慢了，要跑上一两个小时，这不是潘大师该有的表现，于是呆鸟尝试了一下，把处理 500 万条数据的速度压缩到了 10 秒以内，现总结如下：

原始数据分为两列，一列为 id，一列为 date，id 不唯一，每个 id 出现至少一次，并对应不同的 date；date 的排序为由远及近，即日期近的在上面，日期远的在下面，类似一个客户 id 在不同日期进行多次消费

std = pd.Timedelta('4 days')，设定标准间隔天数为 4 天
df['date_shift'] = df.groupby(['id'])['date'].shift(periods=1)，用 groupby 根据 id 分组，以便对每组数据重新排序

注意:每组的第 1 条数据对应序号肯定为 1

从第 2 条数据开始，判断与上一条数据间隔的日期是否为 4 天，怎么判断间隔日期呢？这里就用到了标题里提到的 shift()，将日期数据下移一行

df['date_diff'] = df['date_shift'] - df['date']，两列 date 列的数据相减，计算上下两条记录相隔的天数，注意计算结果的类型为 timedelta
df['date_diff'] = (df['date_diff']-std) \> pd.Timedelta('0 days')，配合 pd.Timedelta() 计算间隔日期是否超过 4 天
df['diff'] = np.where(df['date_diff'] == True, 1,0)，对 pandas 的列数据执行 if-then-else 判断，超过 4 天则为 1，没超过则为 0，新建 “diff” 列储存计算结果，注意：这里的 0 与 1 是为了用于下一步计算
df['diff_cumsum'] = df.groupby(['id'])['diff'].cumsum()，用 cumsum() 为每组数据添加组内序号，注意：这里有一个分组后的计算技巧，这里的累加是针对 diff 列的，也就是上一步计算的 0 与 1，这样计算的目的是实现只有间隔日期超过 4 天序号才递增，没超过 4 天则不递增，计算结果存在 diff_cumsum

这一步虽没什么技术含量，但思路非常重要，就是这种方式，使得计算速度大幅提升。

df['new_id_surfix'] = df['diff_cumsum'] + 1，把 diff_cumsum 列的数据 + 1，即可得到每组的序号，计算结果存入到新 id 的后缀列里即 new_id_surfix
df['new_id'] = df['id'].astype(str) + '-' + df['new_id_surfix'].astype(str)，用 astype(str) 把类型转为字符串，用 + 合并两个，则可生成新的 id