Pandas DataFrame与列表的高效逐元素比较及好友匹配效率优化方案

下面提供两种可直接落地的优化方案，你可以根据实际使用场景选择：

方案1：直接生成你需要的全量布尔矩阵 #

用numpy广播机制可以一次性生成你要的(len(A), len(B))形状的all_bools矩阵，无需循环：

import numpy as np

# 先转成numpy数组方便广播计算
A_arr = np.array(A)
# 提取B的两列好友字段，形状为 (len(B), 2)
friend_cols = B[['friend1', 'friend2']].values

# 广播比较：A_arr[:, None, None] 形状为 (len(A), 1, 1)，和friend_cols比较后生成 (len(A), len(B), 2) 的布尔数组
# 沿最后一维取任意匹配的结果，最终得到形状为 (len(A), len(B)) 的全量布尔矩阵
all_bools = (friend_cols == A_arr[:, None, None]).any(axis=2)

后续使用时直接取对应行即可：my_bool = all_bools[A.index(current_name)]
注意：该方案生成的矩阵内存占用约为750MB（1500050000字节），普通配置的机器都可以正常承载*

方案2：预构建映射字典（内存占用更低、速度更快） #

如果你的使用场景是遍历A中的每个名字做处理，不需要全局随机访问所有布尔行，更推荐用预建索引的方案，内存开销可以忽略不计，查询速度更快：

import pandas as pd

# 把两列好友字段拼接成一个长Series，索引保留B原本的行号
friend_series = pd.concat([B['friend1'], B['friend2']])
# 按好友名字分组，得到「名字 -> 对应B中匹配行的索引列表」的映射字典
friend_to_idx = friend_series.groupby(friend_series).groups

后续使用时直接查询字典即可：

for current_name in A:
    matched_idx = friend_to_idx.get(current_name, [])
    # 直接用matched_idx做后续计算即可，不需要再遍历B
    # 如果确实需要布尔数组，也可以快速生成：
    # my_bool = pd.Series(False, index=B.index)
    # my_bool.loc[matched_idx] = True

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