Интервалы кластеризации

Question

Каждая из строк моего фрейма данных представляет собой интервал, представленный date1 и date2, и идентификатор пользователя. Для каждого идентификатора пользователя мне нужно объединить интервалы, разделенные пробелом ниже определенного порога.

До сих пор для каждого идентификатора пользователя я сортировал строки по дате начала и окончания. Затем я вычисляю пробелы и группы строк на основе этих значений. Затем я добавляю измененные строки в новый фреймворк данных (так я обнаружил, что дезагруппировал DataFrame).

Однако это довольно медленно. Вы видите способы улучшить способ группировки?

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']]))])) 

Answer 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) 

Edit: после комментария OP, я переместил конкатенацию из цикла, чтобы улучшить производительность.

Дальнейшее улучшение может быть не сортировать группы в groupby операции (если есть много пользователей):

grouped = df.groupby('user_id', sort=False) 

или даже группировка вручную путем сортировки df по user_id, а затем добавить условие 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()