如何从DataFrame高效快速构建电影题材关联网络图谱

原代码问题分析 #

• 时间复杂度为O(n²)，10万条数据的情况下会产生100亿次循环，完全无法落地
• 关联逻辑错误：原代码用字符串包含判断r['genres'] in r1['genres']，只能匹配A电影的所有题材字符串完全是B电影题材字符串子集的情况，和你要求的「存在共同题材就关联」的规则完全不符，比如你示例中的1和20就不会被原代码识别为有关联

高效实现思路 #

反过来从题材维度构建关联，而非遍历所有电影两两配对：

1. 先构建「题材-对应所有电影ID」的倒排索引
2. 对每个题材下的所有电影，生成两两配对的边，同一对电影如果在多个题材下重复出现只保留一条边
3. 批量加入图中即可，时间复杂度从O(n²)降到O(所有题材下电影两两配对的总次数)，10万条数据实测几秒就能跑完

优化后代码 #

import networkx as nx
import pandas as pd
from itertools import combinations

# 读取数据
rdf = pd.read_csv("movies.csv")
# 构建图，先加所有节点
g = nx.Graph()
g.add_nodes_from(rdf['movie_id'].tolist())

# 构建题材倒排索引
genre_movies = {}
for _, row in rdf.iterrows():
    mid = row['movie_id']
    # 拆分题材字符串为单个题材列表
    genres = row['genres'].strip().split(',')
    for genre in genres:
        genre = genre.strip()
        if genre not in genre_movies:
            genre_movies[genre] = []
        genre_movies[genre].append(mid)

# 生成所有边，自动去重
edges = set()
for mids in genre_movies.values():
    # 同一个题材下至少2部电影才会产生边
    if len(mids) < 2:
        continue
    # 生成所有两两不重复的组合
    for pair in combinations(mids, 2):
        # 排序后存入set避免(a,b)和(b,a)被识别为不同边
        sorted_pair = tuple(sorted(pair))
        edges.add(sorted_pair)

# 批量加边
g.add_edges_from(edges)
# 导出 adjlist
nx.write_adjlist(g, "movie_graph.csv")

额外性能优化建议 #

如果数据量进一步增大，还可以做如下调整：

• 不需要可视化的话可以不用导入matplotlib，减少资源占用
• 可以用pandas的向量化操作代替iterrows做题材拆分，速度会更快
• 如果边数太多内存不够，可以把边直接写入文件不用全部存在set里，写入前做局部去重即可

内容的提问来源于stack exchange，提问作者abc xyz