如何基于Pandas DataFrame构建类别层级（识别直接父类）

问题说明 #

现有如下Pandas DataFrame：

import pandas as pd

pandas_test=pd.DataFrame(data={'TAGS': [['Category1','Category2','Category3'],
                                        ['Category2','Category4'],
                                        ['Category5','Category4'],
                                        ['Category5','Category4','Category6','Category8'],
                                        ['Category1','Category2'],
                                        ['Category2','Category3']]})

需求规则：同一行中，位于某个类别左侧的紧邻类别即为其直接父类；每行第一个类别的父类为None。需要得到每个类别对应的所有直接父类（部分类别可能存在多个父类）。

现有代码问题 #

用户原有代码如下：

category_counts = {}
for index, row in pandas_test.iterrows():
    categories = row['TAGS']
    for i in range(len(categories)):
        category = categories[i].strip()
        if category not in category_counts:
            category_counts[category] = {'count': 1, 'subcategories': set()}
        else:
            category_counts[category]['count'] += 1
        for j in range(i + 1, len(categories)):
            subcategory = categories[j].strip()
            category_counts[category]['subcategories'].add(subcategory)

# Analyze category_counts dictionary to determine hierarchy
hierarchy = {}
for category, data in category_counts.items():
    subcategories = data['subcategories']
    for subcategory in subcategories:
        if subcategory in category_counts:
            if category not in category_counts[subcategory]['subcategories']:
                hierarchy[subcategory] = category

# Apply hierarchy to categories
for category, parent in hierarchy.items():
    if parent in hierarchy:
        hierarchy[category] = hierarchy[parent]
print(hierarchy)

运行结果：

{'Category3': 'Category1', 'Category2': 'Category1', 'Category4': 'Category5', 'Category6': 'Category5', 'Category8': 'Category5'}

存在的问题：

• Category3的父类应为直接的Category2，而非间接的Category1
• Category6、Category8的父类应为直接的上一级（Category4、Category6），而非间接父类
• Category4存在两个直接父类（Category2和Category5），但结果只保留了一个

解决方案 #

直接遍历每行的类别列表，仅记录每个类别紧邻的左侧父类，用集合存储每个类别的所有父类（自动去重）：

import pandas as pd

pandas_test=pd.DataFrame(data={'TAGS': [['Category1','Category2','Category3'],
                                        ['Category2','Category4'],
                                        ['Category5','Category4'],
                                        ['Category5','Category4','Category6','Category8'],
                                        ['Category1','Category2'],
                                        ['Category2','Category3']]})

# 初始化父类映射，键为类别，值为该类别的所有直接父类集合
parent_map = {}

for tags in pandas_test['TAGS']:
    for idx, category in enumerate(tags):
        # 处理每行第一个类别，父类为None
        if idx == 0:
            if category not in parent_map:
                parent_map[category] = set()
            parent_map[category].add(None)
        else:
            # 取紧邻左侧的类别作为直接父类
            direct_parent = tags[idx-1]
            if category not in parent_map:
                parent_map[category] = set()
            parent_map[category].add(direct_parent)

# 可选：将集合转为列表，提升可读性
parent_map = {k: list(v) for k, v in parent_map.items()}
print(parent_map)

运行结果 #

{
    'Category1': [None],
    'Category2': ['Category1'],
    'Category3': ['Category2'],
    'Category4': ['Category2', 'Category5'],
    'Category5': [None],
    'Category6': ['Category4'],
    'Category8': ['Category6']
}

方案说明 #

• 遍历每行的类别列表时，仅关注当前类别与前一个类别的关系，确保只记录直接父类
• 使用集合存储父类，自动去重同一类别重复出现的相同父类
• 避免了原有代码中递归替换父类导致的层级扁平化问题，同时完整保留了多父类的情况

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