fix: make dfs weight ordering iterative (#38313)

This commit is contained in:
Shuwen Wang
2026-09-08 05:10:22 +00:00
committed by GitHub
parent 91a45ea37e
commit 61d501427d
@@ -266,25 +266,32 @@ def _dfs_weight_order(
node: len(indices) for node, indices in last_node_to_indices.items() node: len(indices) for node, indices in last_node_to_indices.items()
} }
def calc_weight(node: Any) -> None: stack: list[tuple[Any, bool]] = [(root_node, False)]
while stack:
node, visited = stack.pop()
if visited:
weight = node_to_weight.get(node, 0)
for child in node.children.values(): for child in node.children.values():
calc_weight(child) weight += node_to_weight.get(child, 0)
node_to_weight[node] = node_to_weight.get(node, 0) + node_to_weight.get( node_to_weight[node] = weight
child, 0 continue
) stack.append((node, True))
for child in reversed(list(node.children.values())):
calc_weight(root_node) stack.append((child, False))
order: list[int] = [] order: list[int] = []
def append_dfs(node: Any) -> None: stack = [(root_node, False)]
while stack:
node, visited = stack.pop()
if visited:
order.extend(last_node_to_indices.get(node, ()))
continue
children = list(node.children.values()) children = list(node.children.values())
children.sort(key=lambda child: -node_to_weight.get(child, 0)) children.sort(key=lambda child: -node_to_weight.get(child, 0))
for child in children: stack.append((node, True))
append_dfs(child) for child in reversed(children):
order.extend(last_node_to_indices.get(node, ())) stack.append((child, False))
append_dfs(root_node)
return order return order