在Pandas中通过列间值关联判定样本分组的实现方法

给定如下Pandas DataFrame：

rdata = {'id': {0: 'a01', 1: 'a02', 2: 'a03', 3: 'a04', 4: 'a05', 5: 'a06', 6: 'a07', 7: 'a08', 8: 'a09', 9: 'a10', 10: 'a11'}, 'col1': {0: 123456, 1: 1234567, 2: 123456, 3: 123456, 4: 1234, 5: 123456, 6: 123456, 7: 1234, 8: 1234, 9: 12345, 10: 12345}, 'col2': {0: 'a', 1: 'c', 2: 'a', 3: 'a', 4: 'a', 5: 'a', 6: 'a', 7: 'a', 8: 'a', 9: 'b', 10: 'b'}, 'col3': {0: 'AAA', 1: 'BBB', 2: 'CCC', 3: 'DDD', 4: 'EEE', 5: 'EEE', 6: 'EEE', 7: 'EEE', 8: 'EEE', 9: 'GGG', 10: 'GGG'}, 'col4': {0: 555555, 1: 111111, 2: 222222, 3: 222222, 4: 666666, 5: 111111, 6: 111111, 7: 111111, 8: 111111, 9: 444444, 10: 333333}}
df = pd.DataFrame(rdata)

该DataFrame包含11个样本与4个特征列，需依据列间值的关联关系判定样本是否属于同一组，关联规则如下：

• 从col4开始，值为111111的样本（a02、a06、a07、a08、a09）为初始组；
• 111111关联col3的BBB与EEE，而EEE还关联col4的666666，因此将a05纳入主组；
• BBB关联col2的c，EEE关联col2的a，因此将col2中所有值为a和c的样本纳入主组；
• a和c关联col1的123456、1234567、1234，最终a01-a09归为同一组，a10与a11为另一组。

需要实现Python函数same_samples(df)，返回样本分组列表，预期输出示例：

same_samples(df)
# 输出：
[['a01','a02','a03','a04','a05','a06','a07','a08','a09'], 
 ['a10','a11']]

这个问题本质是基于特征值关联的图连通性问题：每个样本是节点，若两个样本通过任意特征列的关联链连通，则属于同一组。可以用**广度优先搜索(BFS)**逐步扩展连通的样本集合。

具体步骤 #

1. 构建特征值到样本的映射：为每个特征列创建字典，记录每个特征值对应的所有样本ID，方便快速查找关联样本；
2. 初始化已访问集合与分组结果：用集合记录已处理的样本避免重复，用列表存储最终分组；
3. 遍历所有样本，对未访问样本启动BFS：从初始种子或任意未访问样本出发，通过特征值关联不断扩展样本集合，直到没有新样本加入；
4. 收集所有连通样本作为一组，直到所有样本都被处理。

代码实现 #

import pandas as pd
from collections import deque

def same_samples(df):
    # 1. 构建各特征值对应的样本ID映射
    feature_maps = {}
    for col in df.columns.drop('id'):
        # 按特征值分组，收集对应的所有样本ID
        col_map = df.groupby(col)['id'].apply(list).to_dict()
        feature_maps[col] = col_map
    
    visited = set()
    groups = []
    
    # 2. 遍历所有样本，处理未访问的样本
    for sample_id in df['id']:
        if sample_id not in visited:
            queue = deque([sample_id])
            visited.add(sample_id)
            current_group = [sample_id]
            
            while queue:
                current_id = queue.popleft()
                # 获取当前样本的所有特征值
                sample_row = df[df['id'] == current_id].iloc[0]
                
                # 遍历每个特征列，找到关联样本
                for col in df.columns.drop('id'):
                    val = sample_row[col]
                    related_samples = feature_maps[col].get(val, [])
                    
                    for s_id in related_samples:
                        if s_id not in visited:
                            visited.add(s_id)
                            current_group.append(s_id)
                            queue.append(s_id)
            
            groups.append(current_group)
    
    return groups

# 测试代码
rdata = {'id': {0: 'a01', 1: 'a02', 2: 'a03', 3: 'a04', 4: 'a05', 5: 'a06', 6: 'a07', 7: 'a08', 8: 'a09', 9: 'a10', 10: 'a11'}, 'col1': {0: 123456, 1: 1234567, 2: 123456, 3: 123456, 4: 1234, 5: 123456, 6: 123456, 7: 1234, 8: 1234, 9: 12345, 10: 12345}, 'col2': {0: 'a', 1: 'c', 2: 'a', 3: 'a', 4: 'a', 5: 'a', 6: 'a', 7: 'a', 8: 'a', 9: 'b', 10: 'b'}, 'col3': {0: 'AAA', 1: 'BBB', 2: 'CCC', 3: 'DDD', 4: 'EEE', 5: 'EEE', 6: 'EEE', 7: 'EEE', 8: 'EEE', 9: 'GGG', 10: 'GGG'}, 'col4': {0: 555555, 1: 111111, 2: 222222, 3: 222222, 4: 666666, 5: 111111, 6: 111111, 7: 111111, 8: 111111, 9: 444444, 10: 333333}}
df = pd.DataFrame(rdata)
print(same_samples(df))

代码说明 #

• 特征映射构建：通过groupby快速收集每个特征值对应的样本ID，后续查找关联样本的时间复杂度接近O(1)；
• BFS扩展：从初始样本出发，不断通过特征值关联找到新样本，确保所有连通样本都被纳入同一组；
• 结果验证：运行代码后会输出符合要求的分组，a01-a09为一组，a10-a11为另一组。

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