Pandas DataFrame嵌套数据多级分组关联值的高效查询方法

原始数据 #

我有如下DataFrame：

数据规则说明：不同的红色取值（第1-3行）被归入“Dark Red”分组，该分组同时隶属于“Red”分组（第7-8行）和“Dark”分组（第14-16行），其中“Red”分组还额外关联了更高层级的“Color”分组（第11-13行）。 #

需求 #

为每个Values列的取值获取所有关联的分组名称列表。

示例 #

• 输入："Scarlet Red"
• 期望输出：['Scarlet Red', 'Dark Red', 'Red', 'Dark', 'Color']

样例数据 #

import pandas as pd
d = {'ID': {0: 'P1',  1: 'C1',  2: 'C2',  3: 'P2',  4: 'C3',  5: 'C4',  6: 'L1',  7: 'P1',  8: 'L2',  9: 'P2',  10: 'LL1',  11: 'L1',  12: 'L2',  13: 'D1',  14: 'P1',  15: 'P2'}, 'Values': {0: 'Dark Red',  1: 'Scarlet Red',  2: 'Cherry Red',  3: 'Dark Blue',  4: 'Steel Blue',  5: 'Coral Blue',  6: 'Red',  7: 'Dark Red',  8: 'Blue',  9: 'Dark Blue',  10: 'Color',  11: 'Red',  12: 'Blue',  13: 'Dark',  14: 'Dark Red',  15: 'Dark Blue'}, 'Type': {0: 'Preferred',  1: 'Name',  2: 'Name',  3: 'Preferred',  4: 'Name',  5: 'Name',  6: 'Preferred',  7: 'Name',  8: 'Preferred',  9: 'Name',  10: 'Preferred',  11: 'Name',  12: 'Name',  13: 'Preferred',  14: 'Name',  15: 'Name'}, 'Group': {0: 1,  1: 1,  2: 1,  3: 2,  4: 2,  5: 2,  6: 3,  7: 3,  8: 4,  9: 4,  10: 5,  11: 5,  12: 5,  13: 6,  14: 6,  15: 6}, 'Level': {0: 0,  1: 0,  2: 0,  3: 0,  4: 0,  5: 0,  6: 1,  7: 1,  8: 1,  9: 1,  10: 2,  11: 2,  12: 2,  13: 1,  14: 1,  15: 1}}
df = pd.DataFrame(d)

当前问题 #

现有实现逻辑如下：

# get the preferred names
df_pref =  df[df['Type'].eq('Preferred')][['Values', 'Group']].rename(columns={'Values': 'Preferred'})
df_merge = df.merge(df_pref, on=['Group'], how='left')

def find_higher_levels(search):
    # search = 'Scarlet Red'
    lst = [search]
    previous_search = None
    while search != previous_search:
        previous_search = search
        search = df_merge[df_merge['Values'].eq(search)]['Preferred'].iloc[-1]
        lst.append(search)
    return lst

find_higher_levels('Scarlet Red')
# Out[85]: ['Scarlet Red', 'Dark Red', 'Dark', 'Dark']

该方案存在两个问题：

1. 输出结果不符合预期，遗漏了Red、Color两个分组
2. 每次查询都需要遍历DataFrame，批量处理所有Values取值时效率低

优化解决方案 #

思路 #

原方案的核心问题是忽略了一个值可能对应多个上层父分组的情况，我们可以先预先构建值到所有直接上层分组的映射字典，再通过广度优先搜索遍历所有关联层级，一次性预计算所有值的结果，后续查询直接取数即可。

实现代码 #

from collections import deque

# 1. 构建值到直接上层父分组的映射
parent_map = {}
# 每个Group对应的Preferred值
group_preferred = df[df['Type'] == 'Preferred'].set_index('Group')['Values'].to_dict()
# 遍历所有行，把每个Value对应的父分组加入映射
for _, row in df.iterrows():
    val = row['Values']
    parent = group_preferred[row['Group']]
    # 排除自己指向自己的情况
    if parent != val:
        if val not in parent_map:
            parent_map[val] = set()
        parent_map[val].add(parent)

# 2. 预计算所有值的关联分组列表
def get_all_groups(val):
    res = set()
    q = deque([val])
    while q:
        current = q.popleft()
        if current not in res:
            res.add(current)
            # 加入所有父节点
            for parent in parent_map.get(current, []):
                q.append(parent)
    # 如果需要和示例一致的顺序，可按每个值的最小Level排序
    return list(res)

# 预计算所有唯一值的结果
all_values = df['Values'].unique()
result_map = {val: get_all_groups(val) for val in all_values}

# 测试
print(result_map['Scarlet Red'])
# 输出：['Scarlet Red', 'Dark Red', 'Red', 'Dark', 'Color'] 顺序可按需求调整

效率说明 #

• 预计算阶段仅需遍历一次全表，加上每个值的层级遍历，总时间复杂度为O(N)，N为数据行数
• 后续查询所有值的关联分组都是O(1)时间，非常适合批量处理场景

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