Pandas: фильтр по дате

Question

У меня есть кадр вроде:

     id         title       date
0  1211  jingle bells 2019-01-15
1  1212  jingle bells 2019-01-15
2  1225      tom boat 2019-06-15
3  2112      tom boat 2019-06-15
4  3122      tom boat 2017-03-15
5  1762      tom boat 2017-03-15

item определяется как группа id с тем же title и date в течение 70 дней с первого дня. Мне нужен словарь из id s, сгруппированный по title, если date находится в пределах 70 дней друг от друга. Ожидаемый результат здесь:

d = {0: [1211,1212], 1: [1225,2112], 2: [3122,1762]}

Любой данный title может иметь неограниченное количество словарных статей или только одну. id являются уникальными для одного заголовка. На данный момент я делаю что-то вроде:

itemlist = []
for i in list(df.title):
    dates = list(df.loc[df.title==i,'date'])
    if (max(dates)-min(dates)).days > 70:
        items = []
        while len(dates)>0:
            extract = [i for i in dates if (i-min(dates)).days<70]
            items.append(list(df.loc[(df.title==i)&(df.date.isin(extract)),'id'])
            dates = [i for i in dates if i not in extract
    else:
        items = [list(df.loc[df.title==i,'id'])]
    itemlist += items
d = {j:i for i in range(len(itemlist)) for j in itemlist[i]}

Это пока не совсем работает, я исправляю ошибки. Тем не менее, я чувствую, что это много итераций - есть идеи, как сделать это лучше? другой приемлемый вывод будет список фреймов данных, по одному на item.

Answer 1

Я думаю, что сортировка вашего фрейма данных может помочь вам решить проблему намного эффективнее.

df = df.sort_values(['title', 'date'])

itemlist = []
counter = 0 # to get items at constant time

for title in set(df.title):
    dates = df.loc[df['title']==title].date.tolist()
    item = []
    min_date = dates[0]
    for date in dates:
        if (date-min_date).days>70: # we need a new item
            itemlist.append(item) # append original item
            item = [df.iloc[counter, 0]] # new item
            min_date = date
        else:
            item.append(df.iloc[counter, 0])
        counter += 1
    itemlist.append(item)

d = {i:j for i,j in enumerate(itemlist)}

print(d)

Несмотря на то, что код стал немного длиннее, есть только два цикла (кроме последнего, который изменяет list на dict), и в общей сложности цикл n_rows, что означает, что он просматривает каждую строку только один раз.
Использование counter заключается в использовании df.iloc, который использует позиционный индекс (вместо меток или условных операторов, таких как df.loc), следовательно, вычисляется быстрее с O (1) .