fix: make dfs weight ordering iterative (#38313)
This commit is contained in:
@@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user