如何用Python递归/迭代法构建指定层级的嵌套树形结构

Got it, let's work through this problem together—you’ve spent two days on this, so let’s break it down clearly and build a solid build_tree function that meets all your requirements.

First, let's recap the core rules we need to follow:

• The hierarchy must start with Year, followed by any combination of Region, Country, Manufacturer, Brand
• The tree is nested: each node has a name (the value of the current hierarchy level), sales (sum of all child node sales, or sum of matching rows if it's a leaf), and optional children (for non-leaf nodes)
• Leaf nodes (last level of the hierarchy) don't have children—their sales is the total of all matching data rows

Step 1: Implement the Function #

We'll use a recursive approach here, since it's perfect for building nested structures. First, we'll add validation for the hierarchy, then a helper function to build each level of the tree.

from itertools import groupby

def build_tree(rows, hierarchy):
    # Validate hierarchy rules first
    if not hierarchy or hierarchy[0] != 'Year':
        raise ValueError("Hierarchy must start with 'Year'")
    
    allowed_levels = {'Year', 'Region', 'Country', 'Manufacturer', 'Brand'}
    for level in hierarchy:
        if level not in allowed_levels:
            raise ValueError(f"Invalid level '{level}'—allowed levels are {allowed_levels}")
    
    # Recursive helper to build each tree level
    def build_level(data, remaining_hierarchy):
        current_level = remaining_hierarchy[0]
        # Sort data first (required for groupby to work correctly)
        sorted_data = sorted(data, key=lambda x: x[current_level])
        nodes = []
        
        for group_key, group_rows in groupby(sorted_data, key=lambda x: x[current_level]):
            group_list = list(group_rows)
            
            if len(remaining_hierarchy) > 1:
                # We have more levels to build—recurse for children
                child_nodes = build_level(group_list, remaining_hierarchy[1:])
                # Roll up sales from children
                total_sales = sum(child['sales'] for child in child_nodes)
                nodes.append({
                    'name': group_key,
                    'sales': total_sales,
                    'children': child_nodes
                })
            else:
                # This is the final level—sum all sales in the group
                total_sales = sum(row['sales'] for row in group_list)
                nodes.append({
                    'name': group_key,
                    'sales': total_sales
                })
        
        return nodes
    
    # Start building from the full hierarchy
    return build_level(rows, hierarchy)

Step 2: Test with Sample Data #

Let's test this with a realistic dataset to make sure it works as expected.

Sample Input #

sales_rows = [
    {'Year': 2023, 'Region': 'North America', 'Country': 'USA', 'Manufacturer': 'Apple', 'Brand': 'iPhone', 'sales': 5000},
    {'Year': 2023, 'Region': 'North America', 'Country': 'USA', 'Manufacturer': 'Apple', 'Brand': 'Mac', 'sales': 3000},
    {'Year': 2023, 'Region': 'Asia', 'Country': 'China', 'Manufacturer': 'Huawei', 'Brand': 'Mate', 'sales': 7000},
    {'Year': 2022, 'Region': 'Europe', 'Country': 'Germany', 'Manufacturer': 'Samsung', 'Brand': 'Galaxy', 'sales': 4000},
    {'Year': 2022, 'Region': 'Europe', 'Country': 'Germany', 'Manufacturer': 'Samsung', 'Brand': 'Z Fold', 'sales': 2000}
]

Test the Hierarchy ['Year', 'Country', 'Manufacturer'] #

tree = build_tree(sales_rows, ['Year', 'Country', 'Manufacturer'])

Expected Output #

[
    {
        "name": 2023,
        "sales": 15000,
        "children": [
            {
                "name": "USA",
                "sales": 8000,
                "children": [
                    {"name": "Apple", "sales": 8000}
                ]
            },
            {
                "name": "China",
                "sales": 7000,
                "children": [
                    {"name": "Huawei", "sales": 7000}
                ]
            }
        ]
    },
    {
        "name": 2022,
        "sales": 6000,
        "children": [
            {
                "name": "Germany",
                "sales": 6000,
                "children": [
                    {"name": "Samsung", "sales": 6000}
                ]
            }
        ]
    }
]

Edge Cases to Keep in Mind #

• Single-level hierarchy: If you pass just ['Year'], the function returns a list of Year nodes, each with their total annual sales and no children.
• Duplicate entries: The groupby and sum logic automatically handles duplicate combinations (like multiple rows for the same Year/Country/Manufacturer).
• Empty input: If rows is empty, the function returns an empty list—no errors (you could add a check for this if needed).

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