python - 在 Python 中查找共同朋友的数量

Question

我有一个用户及其朋友的数据框，如下所示：

user_id | friend_id
1         3
1         4
2         3
2         5
3         4

我想编写一个函数python来计算每对的共同朋友的数量：

user_id | friend_id | num_mutual
1         3           1
1         4           1
2         3           0
2         5           0
3         4           1

目前我有：

def find_mutual(df):
    num_mutual = []
    for i in range(len(df)):
        user, friend = df.loc[i, 'user_id'], df.loc[i, 'friend_id']
        user_list = df[df.user_id == user].friend_id.tolist() + df[df.friend_id == user].user_id.tolist()
        friend_list = df[df.user_id == friend].friend_id.tolist() + df[df.friend_id == friend].user_id.tolist()
        mutual = len(list(set(user_list) & set(friend_list)))
        num_mutual.append(mutual)
    return num_mutual

它适用于小型数据集，但我在具有数百万行的数据集上运行它。运行一切需要很长时间。我知道这不是找到计数的理想方法。Python中有更好的算法吗？提前致谢！

Answer 1

[ugly] 的想法是构建一个以 a 开头user_id并以 same 结尾的 4 点路径user_id。如果存在这样的路径，则 2 个起点有共同的朋友。

我们从：

df
          user_id  friend_id
0        1          3
1        1          4
2        2          3
3        2          5
4        3          4

然后你可以这样做：

dff = df.append(df.rename(columns={"user_id":"friend_id","friend_id":"user_id"}))
df_new = dff.merge(dff, on="friend_id", how="outer")
df_new = df_new[df_new["user_id_x"]!= df_new["user_id_y"]]
df_new = df_new.merge(dff, left_on= "user_id_y", right_on="user_id")
df_new = df_new[df_new["user_id_x"]==df_new["friend_id_y"]]
df_out = df.merge(df_new, left_on=["user_id","friend_id"], right_on=["user_id_x","friend_id_x"], how="left",suffixes=("__","_"))
df_out["count"] = (~df_out["user_id_x"].isnull()).astype(int)
df_out[["user_id__","friend_id","count"]]

   user_id__  friend_id  count
0          1          3      1
1          1          4      1
2          2          3      0
3          2          5      0
4          3          4      1

使用图方法的更优雅和直接的方式

import networkx as nx
g = nx.from_pandas_edgelist(df, "user_id","friend_id")
nx.draw_networkx(g)

然后，您可以将共同朋友的数量标识为存在 3 个节点路径的 2 个相邻节点（2 个朋友）的路径数：

from networkx.algorithms.simple_paths import all_simple_paths
for row in df.itertuples():
    df.at[row[0],"count"] = sum([len(l)==3 for l in list(all_simple_paths(g, row[1], row[2]))])
print(df)
   user_id  friend_id  count
0        1          3    1.0
1        1          4    1.0
2        2          3    0.0
3        2          5    0.0
4        3          4    1.0