如何在DataFrame中匹配两列并按Col1首值分组生成Col3

问题描述 #

给定一个DataFrame，需要比较Col1和Col2列，当记录间存在关联匹配时（例如行2的LG Premier出现在行1的Col2，行3的S Premier出现在行2的Col2），将Col3统一设置为该关联链中Col1的首个值。

示例输入DataFrame：

|Col1        |Col2       |cnt   |
1 |Premier     | LG Premier|3     | 
2 |LG Premier  | S Premier |3     |
3 |K Premier   | S Premier |3     |
4 |Dell        | Dell ABC  |2     |
5 |Dell ABC    | Dell GBC  |2     |

预期输出DataFrame：

|Col1        |Col2       |cnt   |Col3       |
1 |Premier     | LG Premier|3     |Premier    |
2 |LG Premier  | S Premier |3     |Premier    |
3 |K Premier   | S Premier |3     |Premier    |  
4 |Dell        | Dell ABC  |2     |Dell       |
5 |Dell ABC    | Dell GBC  |2     |Dell       |

关联逻辑：行1的Col3设为Premier；行2因LG Premier出现在行1的Col2，Col3设为Premier；行3因S Premier出现在行2的Col2，Col3设为Premier。

解决思路 #

这个问题本质是图的连通分量查找：将Col1和Col2中的每个唯一值视为图的节点，每一行的Col1和Col2之间建立一条边，这样关联的记录就会属于同一个连通分量。之后，为每个连通分量找到最早出现的Col1值，将其作为该分量所有行的Col3值。

我们可以用两种方式实现：使用networkx库快速构建图并查找连通分量，或者手动实现并查集（DSU）结构。

方法一：使用NetworkX库实现 #

import pandas as pd
import networkx as nx

# 构造示例数据
data = {
    'Col1': ['Premier', 'LG Premier', 'K Premier', 'Dell', 'Dell ABC'],
    'Col2': ['LG Premier', 'S Premier', 'S Premier', 'Dell ABC', 'Dell GBC'],
    'cnt': [3, 3, 3, 2, 2]
}
df = pd.DataFrame(data)

# 构建图
G = nx.Graph()
# 添加所有唯一值作为节点
all_values = pd.concat([df['Col1'], df['Col2']]).unique()
G.add_nodes_from(all_values)
# 为每一行的Col1和Col2添加边
for _, row in df.iterrows():
    G.add_edge(row['Col1'], row['Col2'])

# 查找所有连通分量
components = list(nx.connected_components(G))

# 为每个连通分量确定首个Col1值
component_root = {}
for comp in components:
    # 找到该分量中最早出现的Col1对应的行
    mask = df['Col1'].isin(comp)
    first_row_idx = mask[mask].index.min()
    root_val = df.loc[first_row_idx, 'Col1']
    # 将分量内所有值映射到该根值
    for val in comp:
        component_root[val] = root_val

# 生成Col3列
df['Col3'] = df['Col1'].map(component_root)

print(df)

方法二：手动实现并查集（DSU） #

如果不想依赖第三方库，可以手动实现并查集来管理连通分量：

import pandas as pd

class DSU:
    def __init__(self):
        self.parent = {}
    
    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def union(self, x, y):
        if x not in self.parent:
            self.parent[x] = x
        if y not in self.parent:
            self.parent[y] = y
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x != root_y:
            self.parent[root_y] = root_x

# 构造示例数据
data = {
    'Col1': ['Premier', 'LG Premier', 'K Premier', 'Dell', 'Dell ABC'],
    'Col2': ['LG Premier', 'S Premier', 'S Premier', 'Dell ABC', 'Dell GBC'],
    'cnt': [3, 3, 3, 2, 2]
}
df = pd.DataFrame(data)

# 初始化并查集并合并关联节点
dsu = DSU()
for _, row in df.iterrows():
    dsu.union(row['Col1'], row['Col2'])

# 映射每个值到其连通分量的根
all_values = pd.concat([df['Col1'], df['Col2']]).unique()
root_map = {val: dsu.find(val) for val in all_values}

# 为每个根确定对应的首个Col1值
root_to_first_col1 = {}
for root in set(root_map.values()):
    # 获取当前连通分量的所有值
    component_vals = [val for val, r in root_map.items() if r == root]
    # 找到最早出现的Col1行
    mask = df['Col1'].isin(component_vals)
    first_row_idx = mask[mask].index.min()
    root_to_first_col1[root] = df.loc[first_row_idx, 'Col1']

# 生成Col3列
df['Col3'] = df['Col1'].apply(lambda x: root_to_first_col1[root_map[x]])

print(df)

输出结果 #

两种方法都会得到预期的输出：

Col1        Col2  cnt     Col3
0     Premier   LG Premier    3  Premier
1  LG Premier   S Premier    3  Premier
2  K Premier   S Premier    3  Premier
3        Dell    Dell ABC    2     Dell
4    Dell ABC    Dell GBC    2     Dell

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