⚡️ Speed up function find_last_node by 12,973%

[codeflash-ai]

📄 12,973% (129.73x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 44.6 milliseconds → 341 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 130x speedup by eliminating a nested loop antipattern that caused quadratic time complexity.

What changed:
The original code used a nested comprehension: all(e["source"] != n["id"] for e in edges) inside (n for n in nodes ...). This meant for each node, it iterated through ALL edges to check if that node was a source, resulting in O(n × m) complexity where n = number of nodes and m = number of edges.

The optimization precomputes a set of all source IDs once: source_ids = {e["source"] for e in edges}. Then for each node, it performs a single O(1) set membership test: n["id"] not in source_ids. This reduces complexity to O(n + m).

Why it's faster:

• Set lookup is O(1) vs linear scan O(m): The original code performed up to m comparisons per node. The optimized version uses hash-based set membership which averages constant time.
• Single pass over edges: The optimization iterates edges once to build the set, rather than iterating edges n times (once per node).
• Scales dramatically better: Test results show the improvement grows with input size:
  • Small inputs (2-3 nodes): 60-90% faster
  • Medium inputs (10 nodes, 900 edges): 717% faster
  • Large linear chains (500-1000 nodes): 5000-32000% faster

Test case performance:

• Best gains on large linear chains and deeply nested structures where the original O(n×m) penalty compounds
• Moderate gains (50-100% faster) on small graphs with 2-10 nodes
• Slight regression (~15-25% slower) only on trivial empty inputs due to set creation overhead, but these microsecond differences are negligible in practice

The optimization is particularly valuable when this function is called in hot paths processing workflow graphs, DAGs, or dependency trees with hundreds of nodes and edges.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 51 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Click to see Generated Regression Tests

from __future__ import annotations

# imports
import pytest  # used for our unit tests
from src.algorithms.graph import find_last_node

# unit tests

# -------------------- Basic Test Cases --------------------


def test_single_node_no_edges():
    # Single node, no edges: should return the node itself
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 1.00μs (25.0% faster)


def test_two_nodes_one_edge():
    # Two nodes, one edge from 1 to 2: 2 is the last node
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.92μs -> 1.08μs (76.8% faster)


def test_three_nodes_chain():
    # Three nodes in a chain: 1->2->3, so 3 is last node
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.25μs -> 1.25μs (80.0% faster)


def test_multiple_end_nodes():
    # Multiple nodes with no outgoing edges: returns the first such node
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    # Both 2 and 3 have no outgoing edges, should return 2 (first in input order)
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.79μs -> 1.12μs (59.3% faster)


# -------------------- Edge Test Cases --------------------


def test_empty_nodes_and_edges():
    # No nodes, no edges: should return None
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 709ns -> 834ns (15.0% slower)


def test_edges_but_no_nodes():
    # Edges exist but node list is empty: should return None
    nodes = []
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 750ns -> 958ns (21.7% slower)


def test_no_edges_multiple_nodes():
    # Multiple nodes, no edges: returns the first node
    nodes = [{"id": "a"}, {"id": "b"}, {"id": "c"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.21μs -> 958ns (26.1% faster)


def test_node_with_self_loop():
    # Node with a self-loop should not be considered as last node
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 1}]
    # Only node 2 has no outgoing edges
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.12μs (62.9% faster)


def test_all_nodes_have_outgoing_edges():
    # All nodes have outgoing edges: should return None
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.92μs -> 1.25μs (53.3% faster)


def test_node_with_multiple_outgoing_edges():
    # One node with multiple outgoing edges, others are leaves
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}]
    # 2 and 3 are leaves, returns 2 (first in input order)
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.17μs (60.7% faster)


def test_node_ids_are_strings():
    # Node IDs as strings
    nodes = [{"id": "x"}, {"id": "y"}]
    edges = [{"source": "x", "target": "y"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.92μs -> 1.17μs (64.2% faster)


def test_node_id_is_zero():
    # Node ID is zero (falsy value)
    nodes = [{"id": 0}, {"id": 1}]
    edges = [{"source": 1, "target": 0}]
    # Node 0 has no outgoing edges, so returns 0
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 1.04μs (39.9% faster)


def test_duplicate_node_ids():
    # Duplicate node IDs: function should return the first leaf node, even if duplicates exist
    nodes = [{"id": 1}, {"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    # Both node dicts with id 1 have outgoing edges, 2 is the only leaf
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.12μs -> 1.21μs (75.9% faster)


def test_edge_with_nonexistent_source():
    # Edge with source not in nodes: should not affect result
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 3, "target": 1}]
    # Both 1 and 2 have no outgoing edges, returns 1
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 1.08μs (34.6% faster)


def test_edge_with_nonexistent_target():
    # Edge with target not in nodes: should not affect result
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 3}]
    # Only node 2 has no outgoing edges
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.79μs -> 1.08μs (65.5% faster)


def test_nodes_with_extra_fields():
    # Nodes have extra fields, function should not care
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.08μs (69.3% faster)


# -------------------- Large Scale Test Cases --------------------


def test_large_linear_chain():
    # Large chain of 1000 nodes: last node should be returned
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i + 1} for i in range(N - 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.6ms -> 56.6μs (32699% faster)


def test_large_star_graph():
    # One node with edges to 999 others: all leaves, should return first leaf
    N = 1000
    nodes = [{"id": 0}] + [{"id": i} for i in range(1, N)]
    edges = [{"source": 0, "target": i} for i in range(1, N)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 37.8μs -> 19.9μs (89.7% faster)


def test_large_all_nodes_have_outgoing():
    # All nodes have outgoing edges (cycle): should return None
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": (i + 1) % N} for i in range(N)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.5ms -> 56.7μs (32508% faster)


def test_large_multiple_leaves():
    # 10 roots, each with a chain of 99 nodes, total 1000 nodes
    N = 10
    chain_len = 100
    nodes = [{"id": f"{i}-{j}"} for i in range(N) for j in range(chain_len)]
    edges = []
    for i in range(N):
        for j in range(chain_len - 1):
            edges.append({"source": f"{i}-{j}", "target": f"{i}-{j+1}"})
    # There are N leaves: f"{i}-{chain_len-1}" for i in range(N)
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 260μs -> 50.6μs (415% faster)


def test_large_sparse_graph():
    # 1000 nodes, only 10 edges, should return the first node with no outgoing edge
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i + 1} for i in range(10)]
    # First node with no outgoing edge is node 11
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 7.21μs -> 1.75μs (312% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

# imports
import pytest  # used for our unit tests
from src.algorithms.graph import find_last_node

# unit tests


class TestFindLastNodeBasic:
    """Test basic functionality with valid inputs"""

    def test_single_node_no_edges(self):
        """Test with a single node and no edges - node should be the last node"""
        nodes = [{"id": "A"}]
        edges = []
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.29μs -> 1.00μs (29.2% faster)

    def test_two_nodes_linear_chain(self):
        """Test with two nodes in a linear chain A -> B"""
        nodes = [{"id": "A"}, {"id": "B"}]
        edges = [{"source": "A", "target": "B"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.96μs -> 1.12μs (74.0% faster)

    def test_three_nodes_linear_chain(self):
        """Test with three nodes in a linear chain A -> B -> C"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "C"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.38μs -> 1.25μs (90.0% faster)

    def test_branching_structure(self):
        """Test with branching: A -> B, A -> C (B and C are both end nodes)"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = [{"source": "A", "target": "B"}, {"source": "A", "target": "C"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.00μs -> 1.17μs (71.4% faster)


class TestFindLastNodeEmptyInputs:
    """Test edge cases with empty or minimal inputs"""

    def test_empty_nodes_empty_edges(self):
        """Test with both empty lists"""
        nodes = []
        edges = []
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 791ns -> 875ns (9.60% slower)

    def test_empty_nodes_with_edges(self):
        """Test with empty nodes but edges present (invalid state)"""
        nodes = []
        edges = [{"source": "A", "target": "B"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 750ns -> 1.00μs (25.0% slower)

    def test_nodes_with_empty_edges(self):
        """Test with nodes but no edges"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = []
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.21μs -> 959ns (26.1% faster)


class TestFindLastNodeMultipleLastNodes:
    """Test cases where multiple nodes could be considered 'last'"""

    def test_all_nodes_are_last_nodes(self):
        """Test when all nodes have no outgoing edges"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = []
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.25μs -> 875ns (42.9% faster)

    def test_multiple_end_nodes_in_tree(self):
        """Test tree structure with multiple leaf nodes"""
        nodes = [
            {"id": "root"},
            {"id": "left"},
            {"id": "right"},
            {"id": "leaf1"},
            {"id": "leaf2"},
        ]
        edges = [
            {"source": "root", "target": "left"},
            {"source": "root", "target": "right"},
            {"source": "left", "target": "leaf1"},
            {"source": "right", "target": "leaf2"},
        ]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 3.04μs -> 1.33μs (128% faster)


class TestFindLastNodeNoLastNode:
    """Test cases where there is no last node"""

    def test_circular_graph_two_nodes(self):
        """Test with circular dependency: A -> B, B -> A"""
        nodes = [{"id": "A"}, {"id": "B"}]
        edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "A"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.96μs -> 1.17μs (67.8% faster)

    def test_circular_graph_three_nodes(self):
        """Test with circular chain: A -> B -> C -> A"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = [
            {"source": "A", "target": "B"},
            {"source": "B", "target": "C"},
            {"source": "C", "target": "A"},
        ]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.38μs -> 1.25μs (90.0% faster)

    def test_self_referencing_node(self):
        """Test with a node that points to itself"""
        nodes = [{"id": "A"}]
        edges = [{"source": "A", "target": "A"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.46μs -> 1.08μs (34.7% faster)

    def test_all_nodes_have_outgoing_edges(self):
        """Test where every node has at least one outgoing edge"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
        edges = [
            {"source": "A", "target": "B"},
            {"source": "B", "target": "C"},
            {"source": "C", "target": "A"},
        ]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.33μs -> 1.25μs (86.6% faster)


class TestFindLastNodeComplexStructures:
    """Test complex graph structures"""

    def test_disconnected_components(self):
        """Test with disconnected graph components"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}, {"id": "D"}]
        edges = [{"source": "A", "target": "B"}, {"source": "C", "target": "D"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.96μs -> 1.21μs (62.1% faster)

    def test_converging_paths(self):
        """Test multiple paths converging to one node"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}, {"id": "D"}]
        edges = [
            {"source": "A", "target": "D"},
            {"source": "B", "target": "D"},
            {"source": "C", "target": "D"},
        ]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.79μs -> 1.29μs (116% faster)

    def test_node_with_multiple_outgoing_edges(self):
        """Test node with multiple outgoing edges"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}, {"id": "D"}]
        edges = [
            {"source": "A", "target": "B"},
            {"source": "A", "target": "C"},
            {"source": "A", "target": "D"},
        ]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.12μs -> 1.21μs (75.8% faster)

    def test_diamond_structure(self):
        """Test diamond structure: A -> B,C -> D"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}, {"id": "D"}]
        edges = [
            {"source": "A", "target": "B"},
            {"source": "A", "target": "C"},
            {"source": "B", "target": "D"},
            {"source": "C", "target": "D"},
        ]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.88μs -> 1.33μs (116% faster)


class TestFindLastNodeDataIntegrity:
    """Test edge cases related to data integrity issues"""

    def test_edges_reference_nonexistent_nodes(self):
        """Test when edges reference node IDs that don't exist in nodes list"""
        nodes = [{"id": "A"}, {"id": "B"}]
        edges = [
            {"source": "A", "target": "Z"},  # Z doesn't exist
            {"source": "X", "target": "B"},  # X doesn't exist
        ]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.92μs -> 1.21μs (58.7% faster)

    def test_duplicate_node_ids(self):
        """Test with duplicate node IDs in the list"""
        nodes = [{"id": "A"}, {"id": "B"}, {"id": "A"}]
        edges = [{"source": "A", "target": "B"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.79μs -> 1.12μs (59.3% faster)

    def test_nodes_with_additional_properties(self):
        """Test nodes with additional properties beyond 'id'"""
        nodes = [
            {"id": "A", "name": "Node A", "value": 100},
            {"id": "B", "name": "Node B", "value": 200},
        ]
        edges = [{"source": "A", "target": "B"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.88μs -> 1.12μs (66.7% faster)

    def test_edges_with_additional_properties(self):
        """Test edges with additional properties beyond 'source' and 'target'"""
        nodes = [{"id": "A"}, {"id": "B"}]
        edges = [{"source": "A", "target": "B", "weight": 5, "type": "directed"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.79μs -> 1.12μs (59.3% faster)

    def test_numeric_node_ids(self):
        """Test with numeric node IDs instead of strings"""
        nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
        edges = [{"source": 1, "target": 2}, {"source": 2, "target": 3}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.38μs -> 1.25μs (90.0% faster)

    def test_mixed_type_node_ids(self):
        """Test with mixed type node IDs"""
        nodes = [{"id": "A"}, {"id": 2}, {"id": "C"}]
        edges = [{"source": "A", "target": 2}, {"source": 2, "target": "C"}]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 2.54μs -> 1.25μs (103% faster)


class TestFindLastNodeLargeScale:
    """Test performance and scalability with large data samples"""

    def test_large_linear_chain(self):
        """Test with a long linear chain of 500 nodes"""
        # Create nodes: 0 -> 1 -> 2 -> ... -> 499
        num_nodes = 500
        nodes = [{"id": i} for i in range(num_nodes)]
        edges = [{"source": i, "target": i + 1} for i in range(num_nodes - 1)]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 4.63ms -> 28.5μs (16158% faster)

    def test_large_number_of_nodes_few_edges(self):
        """Test with many nodes but few edges"""
        # 800 nodes, only first 10 connected
        num_nodes = 800
        nodes = [{"id": i} for i in range(num_nodes)]
        edges = [{"source": i, "target": i + 1} for i in range(10)]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 6.54μs -> 1.79μs (265% faster)

    def test_wide_branching_structure(self):
        """Test with one root node branching to many leaf nodes"""
        # Root node connects to 500 leaf nodes
        num_leaves = 500
        nodes = [{"id": "root"}] + [{"id": f"leaf_{i}"} for i in range(num_leaves)]
        edges = [{"source": "root", "target": f"leaf_{i}"} for i in range(num_leaves)]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 22.3μs -> 11.1μs (102% faster)

    def test_large_complete_binary_tree(self):
        """Test with a complete binary tree structure"""
        # Create a binary tree with depth 8 (255 nodes)
        nodes = []
        edges = []
        num_nodes = 255  # Complete binary tree with depth 8

        for i in range(num_nodes):
            nodes.append({"id": i})
            left_child = 2 * i + 1
            right_child = 2 * i + 2
            if left_child < num_nodes:
                edges.append({"source": i, "target": left_child})
            if right_child < num_nodes:
                edges.append({"source": i, "target": right_child})

        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 619μs -> 11.2μs (5443% faster)

    def test_large_number_of_edges_same_nodes(self):
        """Test with many edges between the same set of nodes"""
        nodes = [{"id": i} for i in range(10)]
        # Create many edges (multiple edges between same pairs)
        edges = []
        for i in range(9):
            for j in range(100):  # 100 edges from each node to the next
                edges.append({"source": i, "target": i + 1})

        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 167μs -> 20.5μs (717% faster)

    def test_sparse_large_graph(self):
        """Test with large graph where most nodes are disconnected"""
        num_nodes = 1000
        nodes = [{"id": i} for i in range(num_nodes)]
        # Only connect every 10th node
        edges = [{"source": i, "target": i + 10} for i in range(0, num_nodes - 10, 10)]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 6.21μs -> 4.67μs (33.1% faster)

    def test_deeply_nested_structure(self):
        """Test with deeply nested but narrow structure"""
        # Create a structure with 300 levels of nesting
        depth = 300
        nodes = [{"id": f"level_{i}"} for i in range(depth)]
        edges = [
            {"source": f"level_{i}", "target": f"level_{i+1}"} for i in range(depth - 1)
        ]
        codeflash_output = find_last_node(nodes, edges)
        result = codeflash_output  # 1.75ms -> 34.2μs (5004% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-find_last_node-mjs9q8xq and push.