diff --git a/src/algorithms/graph.py b/src/algorithms/graph.py index 777ea3b..95ad690 100644 --- a/src/algorithms/graph.py +++ b/src/algorithms/graph.py @@ -47,7 +47,12 @@ def find_shortest_path(self, start: str, end: str) -> list[str]: def find_last_node(nodes, edges): """This function receives a flow and returns the last node.""" - return next((n for n in nodes if all(e["source"] != n["id"] for e in edges)), None) + # Collect all edge sources into a set for O(1) lookups + edge_sources = {e["source"] for e in edges} + for n in nodes: + if n["id"] not in edge_sources: + return n + return None def find_leaf_nodes(nodes: list[dict], edges: list[dict]) -> list[dict]: