120 lines
4.5 KiB
Python
120 lines
4.5 KiB
Python
def build_tree(records):
|
||
"""
|
||
构建带层级属性的树形结构
|
||
返回: 根节点字典 {id, name, children: [...], level}
|
||
"""
|
||
nodes = {}
|
||
root_candidates = set() # 存储可能的根节点ID [2,7](@ref)
|
||
|
||
# 创建所有节点（暂不设置level）
|
||
for id_str, name, parent_id in records:
|
||
try:
|
||
node_id = int(id_str)
|
||
parent_id_val = int(parent_id) if parent_id and str(parent_id).isdigit() else None
|
||
|
||
nodes[node_id] = {
|
||
'id': node_id,
|
||
'label': name,
|
||
'children': [] # level属性稍后通过递归设置
|
||
}
|
||
|
||
# 标记根节点候选（父节点不存在或无效）
|
||
if parent_id_val is None or parent_id_val not in nodes:
|
||
root_candidates.add(node_id)
|
||
except (ValueError, TypeError):
|
||
continue
|
||
|
||
# 构建父子关系
|
||
for id_str, _, parent_id in records:
|
||
try:
|
||
node_id = int(id_str)
|
||
parent_id_val = int(parent_id) if parent_id and str(parent_id).isdigit() else None
|
||
|
||
# 连接父子节点
|
||
if parent_id_val in nodes and node_id in nodes:
|
||
nodes[parent_id_val]['children'].append(nodes[node_id])
|
||
# 从根候选移除（有父节点说明不是根）
|
||
if node_id in root_candidates:
|
||
root_candidates.remove(node_id)
|
||
except (ValueError, TypeError):
|
||
continue
|
||
|
||
# ============ 新增核心逻辑：递归设置层级属性 ============
|
||
def set_node_level(node, current_level):
|
||
"""递归设置节点层级属性（DFS深度优先遍历）[3,7](@ref)"""
|
||
node['level'] = current_level # 设置当前节点层级
|
||
for child in node['children']:
|
||
set_node_level(child, current_level + 1) # 子节点层级+1
|
||
|
||
# 为所有根节点设置层级
|
||
for root_id in root_candidates:
|
||
set_node_level(nodes[root_id], 0) # 根节点层级从0开始
|
||
|
||
# 返回第一个根节点（或None）
|
||
return nodes[next(iter(root_candidates))] if root_candidates else None
|
||
|
||
|
||
def build_sub_tree(records, target_parent_id=None):
|
||
"""
|
||
构建带层级属性的树形结构
|
||
|
||
参数:
|
||
records: 记录列表，格式为[(id_str, name, parent_id)]
|
||
target_parent_id: 指定的一级标签ID（作为子树根节点）
|
||
|
||
返回:
|
||
带level属性的子树或完整树结构
|
||
"""
|
||
nodes = {}
|
||
root_candidates = set() # 存储可能的根节点ID
|
||
|
||
# 创建所有节点（暂不设置level）
|
||
for id_str, name, parent_id in records:
|
||
try:
|
||
node_id = int(id_str)
|
||
parent_id_val = int(parent_id) if parent_id and str(parent_id).isdigit() else None
|
||
|
||
nodes[node_id] = {
|
||
'id': node_id,
|
||
'label': name,
|
||
'children': [] # level属性稍后通过递归设置
|
||
}
|
||
|
||
if parent_id_val is None or parent_id_val not in nodes:
|
||
root_candidates.add(node_id)
|
||
except (ValueError, TypeError):
|
||
continue
|
||
|
||
# 构建父子关系
|
||
for id_str, _, parent_id in records:
|
||
try:
|
||
node_id = int(id_str)
|
||
parent_id_val = int(parent_id) if parent_id and str(parent_id).isdigit() else None
|
||
|
||
if parent_id_val in nodes and node_id in nodes:
|
||
nodes[parent_id_val]['children'].append(nodes[node_id])
|
||
if node_id in root_candidates:
|
||
root_candidates.remove(node_id)
|
||
except (ValueError, TypeError):
|
||
continue
|
||
|
||
# ============ 新增核心逻辑：递归设置层级属性 ============
|
||
def set_node_level(node, current_level):
|
||
"""递归设置节点层级属性（DFS深度优先遍历）"""
|
||
node['level'] = current_level # 设置当前节点层级
|
||
for child in node['children']:
|
||
set_node_level(child, current_level + 1) # 子节点层级+1[8](@ref)
|
||
|
||
# 场景1：处理子树
|
||
if target_parent_id is not None:
|
||
target_id = int(target_parent_id)
|
||
subtree_root = nodes.get(target_id)
|
||
if subtree_root:
|
||
set_node_level(subtree_root, 1) # 子树根节点从0开始
|
||
return subtree_root
|
||
|
||
# 场景2/3：处理完整树或森林
|
||
for root_id in root_candidates:
|
||
set_node_level(nodes[root_id], 0) # 每个根节点从0开始[7](@ref)
|
||
|
||
return nodes[next(iter(root_candidates))] if root_candidates else nodes |