Pandas. Поиск значений в столбце со списками?

Question

[pintel]

Есть два датафрейма

df1 = pd.DataFrame({'id': [24, 75, 32, 89]})
df2 = pd.DataFrame({'id': [4, 87, 145, 99, 175],
                    'lst': [[24, 56, 78], [24, 32, 89, 54, 127], [67], [78, 89, 34], [12, 45]]})

Как найти все строки в df2, чтобы значения df1['id] были в списках.

Решение через циклы - слишком медленно

df3 = pd.DataFrame()
for i in df2.itertuples():
    for j in df1.itertuples():
        if j[1] in i[2]:
            df3 = df3.append({'id': i[1], 'lst': i[2]}, ignore_index=True)

Метод isin() в данном случае не работает

for j in df1.itertuples():
    df3 = df2[df2['id'].isin(j[1])]

Нужен такой результат

id                    lst
0   4           [24, 56, 78]
1  87  [24, 32, 89, 54, 127]
2  99           [78, 89, 34]

Comments

[Алан Гибизов]

У тебя в df2 id не уникальные - два раза 145.

[pintel]

Алан Гибизов, Эти датафреймы просто, как пример.

[Алан Гибизов]

pintel, ну так одно, когда id уникальны, другое, когда не уникальны. Считаем, что уникальны - переделай пример. Считаем, что не уникальны - напиши об этом отдельное предупреждение.

Answer 1

[Алан Гибизов]

Вот, нацарапал на салфетке:

import pandas as pd
df1 = pd.DataFrame({'id': [24, 75, 32, 89]})
df2 = pd.DataFrame({'id': [4, 87, 145, 99, 146],
                    'lst': [[24, 56, 78], [24, 32, 89, 54, 127], [67], [78, 89, 34], [12, 45]]})

f1 = set(df1['id'])
df22 = df2.set_index('id')

for ident in df22.index:
    f2 = set(df22['lst'][ident])
    result = list(f1&f2)
    if len(result):
        print(ident, list(result))

Comments

[Алан Гибизов]

заменил дублирующийся id на другое значение

[Алан Гибизов]

или надо это получить в новом df?

[pintel]

Алан Гибизов, Нет, спасибо

Answer 2

[Дмитрий]

Как не крутись, но любое решение будет O(n^2), соответсвенно, чтобы это всё отработало максимально быстро, тебе необходимо делать всё при помощи pandas. Но с такой архитектурой df2 у тебя это наврядли получиться, так что думай как преобразовать исходные датафреймы.

Comments

[pintel]

Все-таки нашел решение при помощи Pandas
search это df1
gr_lst это df2

print('len search: ' + str(len(search)))
print('len gr_lst: ' + str(len(gr_lst)))
start = time.process_time()
f1 = set(search)
df22 = gr_lst.set_index('id')
df3 = pd.DataFrame()
for ident in df22.index:
    f2 = set(df22['gr_list'][ident])
    result = f1&f2
    if len(result):
        chunk = pd.DataFrame({'id': ident, 'gr_list': list(result)})
        df3 = df3.append(chunk, ignore_index=True)
print('time: ' + str(time.process_time() - start))
print(df3)

start = time.process_time()
df_search = pd.DataFrame(gr_lst['gr_list'].apply(lambda x: list(set(x) & set(search))))
df_search['id'] = gr_lst['id']
# Удаляем строки с пустыми списками
df_search = df_search[df_search['gr_list'].str.len() != 0]
# Преобразуем столбец со списком значений в строки
df_search = df_search['gr_list'].apply(lambda x: pd.Series(x)).stack().reset_index(level=1, drop=True).to_frame('gr_list').join(df_search[['id']])
# Косметические преобразования
df_search = df_search.astype({'gr_list': int}, copy=False).reset_index()
df_search = df_search[['id', 'gr_list']]
print('time: ' + str(time.process_time() - start))
print(df_search)

В итоге получаем

len search: 3229
len gr_lst: 60578
time: 13.109375
              id    gr_list
0          46774  126842973
1          46774  146897415
2          72269    2507641
3         128191   29046529
4         128191   48437003
...          ...        ...
31178  593502493  195837567
31179  594033097   68291497
31180  594033097   39438402
31181  594033097  126842973
31182  598864135  126842973

[31183 rows x 2 columns]

time: 10.0625
              id    gr_list
0          46774  126842973
1          46774  146897415
2          72269    2507641
3         128191   29046529
4         128191   48437003
...          ...        ...
31178  593502493  195837567
31179  594033097   68291497
31180  594033097   39438402
31181  594033097  126842973
31182  598864135  126842973

[31183 rows x 2 columns]