根据相互联系得到满足条件的值对pandas

Get pairs of values satisfying condition based on mutual connection pandas

假设我有以下数据框:

     index    A      B
     -----------------
      1      A1     B1
      2      A1     B2
      3      A1     B3
      4      A2     B1

我如何编写代码 return 这些对 (Ax,By) 满足这样的条件,即 Ax 与更多不同的 B 连接,而不是 By 与不同的 A 连接。

在这种情况下应该return (A1, B1) 因为A1连接了3个不同的B,但是B1连接了2个不同的A。

这是一种方法(为清楚起见,分几步):

# Drop duplicates in case there are any
df = df.drop_duplicates() 

df["A_count"] = df.groupby("A")["B"].transform("count")
df["B_count"] = df.groupby("B")["A"].transform("count")
df[(df.A_count > df.B_count)]

输出为:

    A   B  A_count  B_count
0  A1  B1        3        2
1  A1  B2        3        1
2  A1  B3        3        1

我们可以将其视为一个图形问题,并检查这些 节点中的哪些 degree 大于 1。然后只需索引两个值都存在的那些行满足条件:

import networkx as nx

G = nx.from_pandas_edgelist(df, source='A', target='B')
keep = [node for node, deg in G.degree() if deg>1]
df[df[['A','B']].isin(keep).all(1)]

   index   A   B
0      1  A1  B1