diff --git a/src/dsa/nodes.py b/src/dsa/nodes.py index 521d24e..e8206fd 100644 --- a/src/dsa/nodes.py +++ b/src/dsa/nodes.py @@ -10,15 +10,10 @@ def find_last_node(nodes, edges): # Function to find all leaf nodes (nodes with no outgoing edges) def find_leaf_nodes(nodes: list[dict], edges: list[dict]) -> list[dict]: - leaf_nodes = [] - for node in nodes: - is_leaf = True - for edge in edges: - if edge["source"] == node["id"]: - is_leaf = False - break - if is_leaf: - leaf_nodes.append(node) + # Build a set of all node ids that have outgoing edges + outgoing_sources = {edge["source"] for edge in edges} + # Nodes whose id is not in outgoing_sources are leaves + leaf_nodes = [node for node in nodes if node["id"] not in outgoing_sources] return leaf_nodes