问题描述
我数据框的每一行都是一个由date1和date2表示的间隔以及一个用户ID。 对于每个用户ID，我需要将这些间隔分组在一起，这些间隔之间的间隔小于某个阈值。
到目前为止，对于每个用户ID，我都按开始日期和结束日期对行进行排序。 然后，我根据这些值计算间隙并分组行。 然后，将修改后的行添加到新的数据框中（这是我发现将数据框取消分组的方式）。
但是，这很慢。 您看到改善我分组方式的方法了吗？
def gap(group):
return group[['date1', 'date2']].min(axis = 1) - \
group.shift()[['date1', 'date2']].max(axis = 1)
def cluster(df, threshold):
df['clusters'] = 0
grouped = df.groupby('user_id')
newdf = pd.DataFrame()
for name, group in grouped:
group = group.sort_values(['date1', 'date2'], ascending = True)
group['gap'] = gap(group)
cuts = group['gap'] > timedelta(threshold)
df2 = group.copy()
for g, d, r in zip(group.loc[cuts, 'gap'], group.loc[cuts, 'date1'], group.loc[cuts, 'date2']):
df2.loc[((df2['date1'] >= d) & (df2['date2'] >= r)), 'clusters'] +=1
df2 = df2.drop('gap', axis = 1)
newdf = pd.concat([newdf, df2])
return newdf
这是它使用的数据的最小样本：
df = pd.DataFrame(dict([('user_id', np.array(['a', 'a', 'a', 'a', 'a', 'a', 'a'])),
('date1', np.array([datetime.strptime(x, "%y%m%d") for x in ['160101', '160103', '160110', '160120', '160130', '160308', '160325']])),
('date2', np.array([datetime.strptime(x, "%y%m%d") for x in ['160107', '160109', '160115', '160126', '160206', '160314', '160402']]))]))
1楼
一个简单的改进是使用cumsum在布尔矢量cuts ：
def cluster2(df, threshold):
df['clusters'] = 0
grouped = df.groupby('user_id')
df_list = []
for name, group in grouped:
group = group.sort_values(['date1', 'date2'], ascending = True)
group['gap'] = gap(group)
print(group)
cuts = group['gap'] > timedelta(threshold)
df2 = group.copy()
df2['clusters'] = cuts.cumsum()
df_list.append(df2)
return pd.concat(df_list)
编辑 ：按照OP的评论，我将串联移出了循环以提高性能。
进一步的改进可能是不对groupby操作中的组进行排序（如果有很多用户）：
grouped = df.groupby('user_id', sort=False)
甚至通过按user_id对df排序，然后添加条件直接在原始数据帧上cuts来手动分组：
df = df.sort_values(['user_id', 'date1', 'date2'], ascending = True)
df['gap'] = gap(df)
cuts = (df['user_id'] != df['user_id'].shift()) | (df['gap'] > timedelta(threshold))
df['clusters'] = cuts.cumsum()