如何为pandas DataFrame中满足指定条件的行分配共同唯一键

现有如下pandas DataFrame：

d = {'id1':['3def','xyx1','roro','p344o','fjfje'],
     'id2':['hhyy','drdr','xyx1','xyx1','xyx1'],
     'score':[100,90,100,90,70]
    }
df1 = pd.DataFrame(d)

原始DataFrame输出：

id1   id2  score
0   3def  hhyy    100
1   xyx1  drdr     90
2   roro  xyx1    100
3  p344o  xyx1     90
4  fjfje  xyx1     70

需求 #

为每一行分配唯一uid，要求满足所有id1和id2存在相等关联且score大于80的行共用同一个uid，期望输出如下：

id1   id2  score  uid
0   3def  hhyy    100    1
1   xyx1  drdr     90    2
2   roro  xyx1    100    2
3  p344o  xyx1     90    2
4  fjfje  xyx1     70    3

现有问题 #

自行编写的代码输出不符合预期：

msk = (((df1['id1'].isin(df1['id2']).shift(-1)) (df1['id2'].isin(df1['id2']).shift(-1))) &\
           (df1['score'] > 80).shift(-1))
df1['uid'] = msk.cumsum()
print(df1)

错误输出：

id1   id2  score  uid
0   3def  hhyy    100    1
1   xyx1  drdr     90    2
2   roro  xyx1    100    3
3  p344o  xyx1     90    3
4  fjfje  xyx1     70    3

这个需求本质是处理id的连通分组问题，使用并查集（Union-Find）可以稳定实现，不会受行顺序影响，完整实现代码如下：

import pandas as pd

# 并查集实现，用于处理id连通关系
class UnionFind:
    def __init__(self):
        self.parent = dict()
    
    def find(self, x):
        if self.parent.get(x, x) != x:
            self.parent[x] = self.find(self.parent[x])
        self.parent.setdefault(x, x)
        return self.parent[x]
    
    def union(self, x, y):
        fx, fy = self.find(x), self.find(y)
        if fx != fy:
            self.parent[fy] = fx

# 构建原始DataFrame
d = {'id1':['3def','xyx1','roro','p344o','fjfje'],
     'id2':['hhyy','drdr','xyx1','xyx1','xyx1'],
     'score':[100,90,100,90,70]
    }
df1 = pd.DataFrame(d)

# 初始化并查集，仅关联score>80的行的id
uf = UnionFind()
high_score_df = df1[df1['score'] > 80]
for _, row in high_score_df.iterrows():
    uf.union(row['id1'], row['id2'])

# 为每个连通组分配唯一uid
group_mapping = dict()
current_uid = 1
# 先处理高分连通组
for _, row in high_score_df.iterrows():
    root = uf.find(row['id1'])
    if root not in group_mapping:
        group_mapping[root] = current_uid
        current_uid += 1

# 生成uid列，低分行单独分配uid
df1['uid'] = df1.apply(
    lambda x: group_mapping[uf.find(x['id1'])] if x['score']>80 else current_uid + x.name,
    axis=1
)
# 修正uid为连续编号
df1['uid'] = pd.factorize(df1['uid'])[0] + 1

print(df1)

运行后输出完全符合预期：

id1   id2  score  uid
0   3def  hhyy    100    1
1   xyx1  drdr     90    2
2   roro  xyx1    100    2
3  p344o  xyx1     90    2
4  fjfje  xyx1     70    3

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