Python脚本实现列表数据转层级结构的需求及代码修复

需求说明 #

• 原始列表的第一个数据为根节点
• 每条数据的user_id对应父节点的id，匹配成功则作为该父节点的子节点
• 支持任意复杂的多层级结构

原始数据 #

[
    {
        "user_id": 1164146,
        "id": 1158720,
        "ancestors": ",8,77169,1164146,",
        "<OTHER DATA>": ""
    },
    {
        "user_id": 1158720,
        "id": 1259377,
        "ancestors": ",8,77169,1164146,1158720,",
        "<OTHER DATA>": ""
    },
    {
        "user_id": 1259377,
        "id": 1259378,
        "ancestors": ",8,77169,1164146,1158720,1259377,",
        "<OTHER DATA>": ""
    },
    {
        "user_id": 1158720,
        "id": 1259379,
        "ancestors": ",8,77169,1164146,1158720,",
        "<OTHER DATA>": ""
    }
]

期望层级结构 #

[
    {
        "user_id": 1164146,
        "id": 1158720,
        "ancestors": ",8,77169,1164146,",
        "<OTHER DATA>": "",
        "children": [
            {
                "user_id": 1158720,
                "id": 1259377,
                "ancestors": ",8,77169,1164146,1158720,",
                "<OTHER DATA>": "",
                "children": [
                    {
                        "user_id": 1259377,
                        "id": 1259378,
                        "ancestors": ",8,77169,1164146,1158720,1259377,",
                        "<OTHER DATA>": "",
                        "children": []
                    }
                ]
            },
            {
                "user_id": 1158720,
                "id": 1259379,
                "ancestors": ",8,77169,1164146,1158720,",
                "<OTHER DATA>": "",
                "children": []
            }
        ]
    }
]

原代码问题分析 #

原代码仅在顶层列表中查找父节点，无法遍历深层嵌套的子节点，导致多层级的节点会错误挂载到根节点下，无法形成正确的层级结构。

修正后的代码 #

def create_nested_structure(data):
    if not data:
        return []
    
    # 用字典存储所有节点，key为id，实现快速父节点查找
    node_map = {}
    # 初始化所有节点并添加children字段
    for item in data:
        node = item.copy()
        node["children"] = []
        node_map[node["id"]] = node
    
    # 构建层级关系
    root = node_map[data[0]["id"]]
    for node in node_map.values():
        user_id = node["user_id"]
        # 找到父节点且当前节点不是根节点时，挂载到父节点的children下
        if user_id in node_map and node != root:
            node_map[user_id]["children"].append(node)
    
    return [root]

代码说明 #

1. 节点映射表：通过node_map字典以节点id为键存储所有节点，将父节点查找的时间复杂度降为O(1)，大幅提升效率。
2. 初始化节点：提前为每个节点添加空的children字段，避免后续重复判断。
3. 构建层级：遍历所有节点，通过user_id在映射表中找到对应父节点，将当前节点添加到父节点的children列表中。
4. 返回结果：返回包含根节点的列表，根节点为原始数据的第一个元素。

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