⚡️ Speed up function find_last_node by 9,272%

[codeflash-ai]

📄 9,272% (92.72x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 11.5 milliseconds → 122 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 93x speedup (9272% faster) by eliminating redundant iterations over the edges collection.

Key Optimization:

The original code uses a nested generator expression that iterates over edges for every node being checked:

next((n for n in nodes if all(e["source"] != n["id"] for e in edges)), None)

This creates O(N × M) comparisons where N = number of nodes and M = number of edges. For 500 nodes and 499 edges, this means ~250,000 comparisons.

The optimized version:

1. Detects re-iterable containers (lists, tuples) vs single-use iterators
2. For re-iterable edges (the common case): Builds a set of all source IDs upfront in O(M) time, then checks membership in O(1) per node, reducing to O(N + M) total complexity
3. For iterator edges (rare case): Preserves original single-pass semantics to avoid breaking behavior when edges can only be consumed once

Performance Impact:

The test results show massive improvements on larger datasets:

• test_large_scale_last_node_at_end_of_long_list: 17,559% faster (4.39ms → 24.9μs)
• test_find_last_node_many_nodes_many_edges: 11,880% faster (2.07ms → 17.3μs)
• test_find_last_node_deep_chain_performance: 17,678% faster (4.41ms → 24.8μs)

Even small graphs benefit significantly (50-120% faster) due to avoiding the nested iteration overhead. The optimization trades a tiny bit of upfront work (building the set) for dramatically reduced per-node checking cost, especially beneficial when nodes >> edges or when many nodes need checking before finding a match.

Edge cases like empty inputs show minimal regression (<10% slower) because the set construction overhead dominates when there's no work to do, but this is an acceptable tradeoff given the dramatic wins on realistic workloads.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

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


def test_basic_single_last_node():
    # Basic scenario: one node is the target of an edge and the other is not a source -> should return the latter.
    nodes = [{"id": "A"}, {"id": "B"}]  # two nodes in the flow
    edges = [
        {"source": "A", "target": "B"}
    ]  # A is a source, B is not a source anywhere
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.58μs -> 875ns (80.9% faster)


def test_basic_multiple_candidates_returns_first_in_nodes_order():
    # If multiple nodes qualify as "last" (none are sources), the function should return the first such node
    nodes = [{"id": "X"}, {"id": "Y"}, {"id": "Z"}]  # multiple nodes that could be last
    edges = [
        {"source": "not_in_nodes"}
    ]  # no edge source matches any node id -> all nodes qualify
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.21μs -> 750ns (61.1% faster)


def test_empty_nodes_returns_none():
    # Edge case: no nodes provided -> there is no "last node", should return None
    nodes = []  # empty nodes list
    edges = [{"source": 1}]  # edges exist but there are no nodes to evaluate
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 625ns -> 666ns (6.16% slower)


def test_no_edges_returns_first_node():
    # Edge case: empty edges list -> every node is not a source, so the first node should be returned
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = []  # no edges, so all() over empty iterable is True for every node
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.00μs -> 667ns (49.9% faster)


def test_no_candidate_returns_none_when_all_nodes_are_sources():
    # Edge case: every node appears as a source in at least one edge -> function should return None
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1}, {"source": 2}]  # both nodes are sources
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.54μs -> 833ns (85.1% faster)


def test_type_sensitive_matching_between_strings_and_ints():
    # Ensure the function uses strict equality: string "1" is not equal to integer 1
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": "1"}]  # string "1" does not match int 1
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.29μs -> 791ns (63.2% faster)


def test_handles_none_id_values_correctly():
    # Nodes and edges may contain None as an id/source; equality should still work
    nodes = [{"id": None}, {"id": "A"}]
    edges = [{"source": "A"}]  # "A" is a source, None is not a source here
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.29μs -> 750ns (72.1% faster)


def test_malformed_edge_missing_source_key_raises_keyerror():
    # If an edge dict lacks the "source" key, accessing e["source"] should raise KeyError
    nodes = [{"id": "A"}]
    edges = [{"src": "A"}]  # intentionally malformed edge (missing "source")
    # The implementation directly accesses e["source"], so KeyError is the expected outcome
    with pytest.raises(KeyError):
        codeflash_output = find_last_node(nodes, edges)
        _ = codeflash_output  # 1.46μs -> 875ns (66.6% faster)


def test_edges_with_additional_keys_do_not_affect_result():
    # Edges may contain extra keys; only "source" is relevant for the function
    nodes = [{"id": "a"}, {"id": "b"}]
    edges = [
        {"source": "a", "weight": 10, "meta": "x"},
        {"source": "c", "weight": 5, "meta": "y"},
    ]  # only a is a source among nodes; b is not
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.67μs -> 917ns (81.8% faster)


def test_duplicate_node_ids_behavior():
    # When node ids are duplicated, each node dict is still evaluated; if the id appears as a source,
    # all nodes with that id should be considered sources (because equality is id-based).
    nodes = [{"id": "dup"}, {"id": "dup"}, {"id": "unique"}]
    edges = [{"source": "dup"}]  # both first nodes have id "dup" which is a source
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.79μs -> 875ns (105% faster)


def test_nodes_order_matters_when_candidate_is_not_first():
    # Ensure that if the first node is a source and a later node is not, the later node is returned
    nodes = [{"id": "n1"}, {"id": "n2"}, {"id": "n3"}]
    edges = [
        {"source": "n1"},
        {"source": "n3"},
    ]  # n1 and n3 are sources; n2 is the only non-source
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.67μs -> 916ns (82.0% faster)


def test_large_scale_last_node_at_end_of_long_list():
    # Large-scale: create up to 500 nodes and edges so function must iterate through many items.
    # Constraint enforced: avoid loops >1000 steps; 500 is well under that.
    n = 500  # number of nodes (kept <1000 per instructions)
    # Create nodes with ids 0..499 as dictionaries
    nodes = [{"id": i} for i in range(n)]
    # Create edges that mark nodes 0..n-2 as sources; node n-1 will not be a source
    edges = [{"source": i, "target": i + 1} for i in range(n - 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 4.39ms -> 24.9μs (17559% faster)


def test_large_scale_many_candidates_first_is_returned():
    # Large-scale scenario where many nodes qualify; ensure the very first qualifying node is returned.
    # Build 300 nodes where nodes with even ids are not sources; the first node in nodes list that is not
    # a source should be returned (node with id 0 in this setup).
    n = 300
    nodes = [{"id": i} for i in range(n)]
    # Make all odd ids sources; even ids are not listed as sources
    edges = [{"source": i} for i in range(1, n, 2)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 7.04μs -> 4.67μs (50.9% faster)


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

import pytest
from src.algorithms.graph import find_last_node


def test_find_last_node_single_node_no_edges():
    """Test with a single node and no edges - the only node is the last node."""
    nodes = [{"id": 1, "label": "A"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.00μs -> 667ns (49.9% faster)


def test_find_last_node_linear_chain():
    """Test with a linear chain of nodes: A -> B -> C. C is the last node."""
    nodes = [{"id": 1, "label": "A"}, {"id": 2, "label": "B"}, {"id": 3, "label": "C"}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 958ns (95.7% faster)


def test_find_last_node_two_nodes_one_edge():
    """Test with two nodes where one is the source. The non-source is the last."""
    nodes = [{"id": 1, "label": "Start"}, {"id": 2, "label": "End"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 834ns (79.9% faster)


def test_find_last_node_diamond_shape():
    """Test with diamond-shaped graph: A -> B,C and both -> D. D is the last node."""
    nodes = [
        {"id": 1, "label": "A"},
        {"id": 2, "label": "B"},
        {"id": 3, "label": "C"},
        {"id": 4, "label": "D"},
    ]
    edges = [
        {"source": 1, "target": 2},
        {"source": 1, "target": 3},
        {"source": 2, "target": 4},
        {"source": 3, "target": 4},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.29μs -> 1.04μs (120% faster)


def test_find_last_node_with_extra_attributes():
    """Test that nodes with extra attributes still work correctly."""
    nodes = [
        {"id": 1, "label": "Start", "color": "red", "size": 10},
        {"id": 2, "label": "End", "color": "blue", "size": 20},
    ]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.38μs -> 833ns (65.1% faster)


def test_find_last_node_multiple_sources_one_target():
    """Test with multiple nodes as sources pointing to one target."""
    nodes = [
        {"id": 1, "label": "A"},
        {"id": 2, "label": "B"},
        {"id": 3, "label": "C"},
        {"id": 4, "label": "Final"},
    ]
    edges = [
        {"source": 1, "target": 4},
        {"source": 2, "target": 4},
        {"source": 3, "target": 4},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.17μs -> 1.00μs (117% faster)


def test_find_last_node_empty_nodes_list():
    """Test with an empty nodes list - should return None."""
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 666ns -> 542ns (22.9% faster)


def test_find_last_node_empty_edges_list():
    """Test with empty edges - first node is the last node."""
    nodes = [{"id": 1, "label": "Only"}, {"id": 2, "label": "Node"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.00μs -> 625ns (60.0% faster)


def test_find_last_node_all_nodes_are_sources():
    """Test where every node is a source of some edge - no last node."""
    nodes = [{"id": 1, "label": "A"}, {"id": 2, "label": "B"}, {"id": 3, "label": "C"}]
    edges = [
        {"source": 1, "target": 2},
        {"source": 2, "target": 3},
        {"source": 3, "target": 1},  # Creates a cycle
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.79μs -> 917ns (95.4% faster)


def test_find_last_node_single_self_loop():
    """Test with a node that has a self-loop - is itself a source."""
    nodes = [{"id": 1, "label": "Loop"}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.12μs -> 750ns (50.0% faster)


def test_find_last_node_with_missing_target_in_nodes():
    """Test where edge target doesn't exist in nodes - source still excluded."""
    nodes = [{"id": 1, "label": "A"}, {"id": 2, "label": "B"}]
    edges = [{"source": 1, "target": 99}]  # Target 99 doesn't exist in nodes
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 875ns (66.6% faster)


def test_find_last_node_node_ids_as_strings():
    """Test with string-based node IDs."""
    nodes = [
        {"id": "node_a", "label": "A"},
        {"id": "node_b", "label": "B"},
        {"id": "node_c", "label": "C"},
    ]
    edges = [
        {"source": "node_a", "target": "node_b"},
        {"source": "node_b", "target": "node_c"},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 958ns (104% faster)


def test_find_last_node_duplicate_edges():
    """Test with duplicate edges - should still find last node correctly."""
    nodes = [{"id": 1, "label": "Start"}, {"id": 2, "label": "End"}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 2}]  # Duplicate edge
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.54μs -> 958ns (61.0% faster)


def test_find_last_node_node_with_only_id():
    """Test with minimal node structure - only id field."""
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.42μs -> 833ns (70.1% faster)


def test_find_last_node_large_id_numbers():
    """Test with very large node ID numbers."""
    nodes = [{"id": 999999999, "label": "A"}, {"id": 1000000000, "label": "B"}]
    edges = [{"source": 999999999, "target": 1000000000}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.54μs -> 833ns (85.0% faster)


def test_find_last_node_negative_ids():
    """Test with negative node IDs."""
    nodes = [{"id": -1, "label": "A"}, {"id": -2, "label": "B"}]
    edges = [{"source": -1, "target": -2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 1.00μs (45.8% faster)


def test_find_last_node_zero_id():
    """Test with zero as a node ID."""
    nodes = [{"id": 0, "label": "Zero"}, {"id": 1, "label": "One"}]
    edges = [{"source": 0, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 834ns (79.9% faster)


def test_find_last_node_complex_node_structure():
    """Test with nodes containing complex nested data."""
    nodes = [
        {"id": 1, "label": "A", "meta": {"color": "red", "tags": ["start"]}},
        {"id": 2, "label": "B", "meta": {"color": "blue", "tags": ["end"]}},
    ]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.42μs -> 833ns (70.1% faster)


def test_find_last_node_large_linear_chain():
    """Test with a large linear chain of nodes (100 nodes)."""
    # Create a chain: 0 -> 1 -> 2 -> ... -> 99
    num_nodes = 100
    nodes = [{"id": i, "label": f"Node_{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  # 203μs -> 6.50μs (3037% faster)


def test_find_last_node_large_branching_tree():
    """Test with a large tree structure where many nodes branch to one final node."""
    # Create nodes 0, 1, 2, ..., 99 all pointing to node 100
    num_sources = 100
    nodes = [{"id": i, "label": f"Node_{i}"} for i in range(num_sources + 1)]
    edges = [{"source": i, "target": num_sources} for i in range(num_sources)]

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 204μs -> 6.38μs (3114% faster)


def test_find_last_node_large_complete_subgraph():
    """Test with a large complete subgraph where all nodes connect to a final node."""
    # Create a complete graph where nodes 0-49 all have edges, then point to node 50
    num_nodes_before_final = 50
    final_id = 100
    nodes = [{"id": i, "label": f"Node_{i}"} for i in range(num_nodes_before_final)]
    nodes.append({"id": final_id, "label": "Final"})

    # Create edges between first 50 nodes
    edges = []
    for i in range(num_nodes_before_final - 1):
        edges.append({"source": i, "target": i + 1})
    # Final edges point to the final node
    for i in range(num_nodes_before_final - 1, num_nodes_before_final):
        edges.append({"source": i, "target": final_id})

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 58.8μs -> 3.75μs (1469% faster)


def test_find_last_node_many_nodes_many_edges():
    """Test with 200 nodes and 500 edges in a complex DAG structure."""
    num_nodes = 200
    nodes = [{"id": i, "label": f"Node_{i}"} for i in range(num_nodes)]

    # Create edges that form a DAG with last node (199) having no outgoing edges
    edges = []
    # Layer-wise connections with skip connections
    for i in range(num_nodes - 1):
        # Connect each node to the next few nodes
        for j in range(i + 1, min(i + 4, num_nodes)):
            if j < num_nodes - 1:  # Don't add edges from the last node
                edges.append({"source": i, "target": j})

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.07ms -> 17.3μs (11880% faster)


def test_find_last_node_many_isolated_nodes_one_chain():
    """Test with many isolated nodes and one chain - only the chain's end is found."""
    # Create 50 isolated nodes and 1 chain of 50 nodes
    isolated_nodes = [{"id": i, "label": f"Isolated_{i}"} for i in range(50)]
    chain_nodes = [{"id": 100 + i, "label": f"Chain_{i}"} for i in range(50)]
    nodes = isolated_nodes + chain_nodes

    edges = [{"source": 100 + i, "target": 100 + i + 1} for i in range(49)]

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 3.12μs -> 2.12μs (47.1% faster)


def test_find_last_node_performance_many_edges_one_source():
    """Test performance with one node as source pointing to many other nodes."""
    # Node 0 points to nodes 1-99
    num_nodes = 100
    nodes = [{"id": i, "label": f"Node_{i}"} for i in range(num_nodes)]
    edges = [{"source": 0, "target": i} for i in range(1, num_nodes)]

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 5.33μs -> 2.67μs (100% faster)


def test_find_last_node_wide_graph_structure():
    """Test with a wide graph where many nodes branch but don't connect to final."""
    # Create 100 nodes where nodes 0-49 are sources and 50-99 are isolated
    num_nodes = 100
    nodes = [{"id": i, "label": f"Node_{i}"} for i in range(num_nodes)]

    # Nodes 0-48 point to node 49
    edges = [{"source": i, "target": 49} for i in range(49)]

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 56.6μs -> 3.67μs (1444% faster)


def test_find_last_node_deep_chain_performance():
    """Test with a very deep chain to ensure no performance degradation."""
    # Create a chain: 0 -> 1 -> 2 -> ... -> 500
    num_nodes = 500
    nodes = [{"id": i, "label": f"Node_{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.41ms -> 24.8μs (17678% faster)


def test_find_last_node_multi_level_diamond():
    """Test with multiple nested diamond patterns."""
    # Create pattern: A -> B,C -> D -> E,F -> G
    nodes = [
        {"id": 1, "label": "A"},
        {"id": 2, "label": "B"},
        {"id": 3, "label": "C"},
        {"id": 4, "label": "D"},
        {"id": 5, "label": "E"},
        {"id": 6, "label": "F"},
        {"id": 7, "label": "G"},
    ]
    edges = [
        {"source": 1, "target": 2},
        {"source": 1, "target": 3},
        {"source": 2, "target": 4},
        {"source": 3, "target": 4},
        {"source": 4, "target": 5},
        {"source": 4, "target": 6},
        {"source": 5, "target": 7},
        {"source": 6, "target": 7},
    ]

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 3.79μs -> 1.33μs (184% 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-mlddarcf and push.