如何从DataFrame列表列提取排除自身的Top3关联项？

问题背景 #

原始DataFrame如下：

需求是提取每个item，并为其生成排除自身的Top3关联项，预期结果：

用户原有代码未实现排除自身的逻辑，导致输出错误，原有代码：

data = [['icecream', [['A', 0.9],['B', 0.6],['C',0.5],['D',0.3],['E',0.1]]], 
        ['potato', [['U', 0.8],['V', 0.7],['W',0.4],['X',0.3],['Y',0.2]]]]

df = pd.DataFrame(data, columns=['tag', 'list'])
df

--

temp = {}
for idx, row in df.iterrows():
    for item in row["list"]:
        temp[item[0]] = row["tag"]

top_items = {}
for idx, row in df.iterrows():
    top_items[row["tag"]] = row["list"]

similar = []
for item, category in temp.items():
    top_3 = top_items.get(category)
    sample = top_3[:3]
    similar.append([item, sample])

df = pd.DataFrame(similar)
df.columns = ["item", "top_3"]

修正后的代码 #

核心是在生成top_3时，先过滤掉当前item，再取前3项：

import pandas as pd

data = [['icecream', [['A', 0.9],['B', 0.6],['C',0.5],['D',0.3],['E',0.1]]], 
        ['potato', [['U', 0.8],['V', 0.7],['W',0.4],['X',0.3],['Y',0.2]]]]

df = pd.DataFrame(data, columns=['tag', 'list'])

# 建立item到tag的映射
item_to_tag = {}
for _, row in df.iterrows():
    for sub_item in row["list"]:
        item_to_tag[sub_item[0]] = row["tag"]

# 建立tag到对应列表的映射
tag_to_list = df.set_index('tag')['list'].to_dict()

similar = []
for item, tag in item_to_tag.items():
    # 过滤掉当前item，然后取前3
    filtered_list = [sub for sub in tag_to_list[tag] if sub[0] != item]
    top_3 = filtered_list[:3]
    similar.append([item, top_3])

result_df = pd.DataFrame(similar, columns=["item", "top_3"])
print(result_df)

更优的pandas原生方法（避免循环） #

利用explode展开列表，再通过merge关联同tag下的其他项，最后分组聚合取Top3：

import pandas as pd

data = [['icecream', [['A', 0.9],['B', 0.6],['C',0.5],['D',0.3],['E',0.1]]], 
        ['potato', [['U', 0.8],['V', 0.7],['W',0.4],['X',0.3],['Y',0.2]]]]

df = pd.DataFrame(data, columns=['tag', 'list'])

# 展开列表为单独行，拆分item和score
df_expanded = df.explode('list').reset_index(drop=True)
df_expanded[['item', 'score']] = pd.DataFrame(df_expanded['list'].tolist(), index=df_expanded.index)

# 同tag内关联所有其他项
merged = df_expanded.merge(df_expanded, on='tag', suffixes=('_self', '_other'))

# 排除自身，按item_self分组，取score_other降序的前3项
merged = merged[merged['item_self'] != merged['item_other']]
top3 = merged.groupby('item_self').apply(
    lambda x: x.sort_values('score_other', ascending=False).head(3)[['item_other', 'score_other']].values.tolist()
).reset_index(name='top_3')

print(top3)

这个方法更符合pandas的向量化操作思路，数据量较大时效率更高。

内容的提问来源于stack exchange，提问作者trojan horse