⚡️ Speed up function find_last_node by 6,295%

[codeflash-ai]

📄 6,295% (62.95x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 9.89 milliseconds → 155 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 6294% speedup (from 9.89ms to 155μs) by eliminating nested iteration through a fundamental algorithmic improvement.

Key Optimization

Original approach: For each node, the code uses all(e["source"] != n["id"] for e in edges) which requires checking every edge for every node - an O(n×m) operation where n is the number of nodes and m is the number of edges.

Optimized approach: Pre-computes a set of all source IDs once (sources = {e["source"] for e in edges}), then performs O(1) membership checks (n["id"] not in sources). This reduces complexity to O(n+m).

Why This Is Faster

1. Set lookup vs repeated iteration: Python set membership tests are O(1) on average, while iterating through all edges for each node is O(m). With multiple nodes, the original approach repeats this m-times traversal n times.

2. Single pass through edges: Building the sources set requires one pass through edges (38.7% of optimized runtime), but this one-time cost is amortized across all node checks.

3. Type guard optimization: The isinstance() check ensures we only use the optimized path for reusable collections, falling back to the original logic for iterators that can't be traversed multiple times.

Performance Characteristics

The optimization excels when:

• Large graphs: The test with 700 nodes shows 24,935% speedup (8.77ms → 35.0μs)
• Dense connections: More edges mean greater savings from avoiding repeated iteration
• Multiple candidates: When checking many nodes, the amortized cost of building the set pays off significantly

The speedup grows dramatically with scale - from ~12-62% on tiny graphs (2-3 nodes) to 1600-2100% on graphs with 100+ nodes, demonstrating the algorithmic improvement's impact on real workloads.

Edge Case Handling

The if not sources: check preserves the original behavior where empty edges means no node has outgoing edges, so the first node is returned without accessing n["id"] (avoiding potential KeyError on malformed data).

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

from __future__ import annotations

import copy  # used to verify inputs are not mutated

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


def test_basic_chain_last_node():
    # Basic scenario: a simple linear flow 1 -> 2 -> 3
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]  # ordered nodes
    edges = [{"source": 1}, {"source": 2}]  # outgoing from 1 and 2
    # the last node (no outgoing edge from id 3) should be the dict with id 3
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.04μs -> 1.42μs (44.0% faster)


def test_multiple_sinks_returns_first_sink_in_order():
    # If several nodes have no outgoing edges, the function should return the first such node
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A"}]  # only A has outgoing edge, B and C are sinks
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.58μs -> 1.33μs (18.8% faster)


def test_no_nodes_returns_none():
    # Edge case: empty sequence of nodes -> nothing to return
    nodes = []
    edges = [{"source": 1}, {"source": 2}]
    codeflash_output = find_last_node(nodes, edges)  # 667ns -> 1.17μs (42.8% slower)


def test_no_edges_returns_first_node():
    # Edge case: no edges means every node has no outgoing edges, so return the first node
    nodes = [{"id": 10}, {"id": 20}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 1.04μs -> 1.00μs (4.20% faster)


def test_edges_missing_source_key_raises_keyerror():
    # If an edge dictionary does not contain 'source', accessing e["source"] should raise KeyError
    nodes = [{"id": 1}]
    edges = [{"src": 1}]  # wrong key
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 1.54μs -> 958ns (61.0% faster)


def test_nodes_missing_id_with_nonempty_edges_raises_keyerror():
    # If nodes do not contain 'id' and edges are non-empty, accessing n["id"] should raise KeyError
    nodes = [{"identifier": 1}]  # wrong key for node id
    edges = [{"source": 1}]  # non-empty so n["id"] will be accessed
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 1.54μs -> 1.54μs (0.000% faster)


def test_nodes_non_dict_type_raises_typeerror_when_edges_nonempty():
    # If nodes are not subscriptable by "id" (e.g., integers), attempting n["id"] raises TypeError
    nodes = [1]  # integer node, not a mapping
    edges = [{"source": 1}]  # non-empty triggers evaluation of n["id"]
    with pytest.raises(TypeError):
        find_last_node(nodes, edges)  # 1.88μs -> 1.75μs (7.14% faster)


def test_duplicate_ids_all_have_outgoing_results_in_none():
    # When multiple nodes share the same id and that id appears as a source,
    # none of those nodes should be considered last -> return None
    nodes = [{"id": 1}, {"id": 1}]
    edges = [{"source": 1}]
    codeflash_output = find_last_node(nodes, edges)  # 1.54μs -> 1.29μs (19.3% faster)


def test_input_sequences_not_mutated():
    # The function should not mutate the input structures (nodes, edges)
    nodes = [{"id": "x"}, {"id": "y"}]
    edges = [{"source": "x"}]
    # deep copy the inputs to compare after function call
    nodes_copy = copy.deepcopy(nodes)
    edges_copy = copy.deepcopy(edges)
    codeflash_output = find_last_node(nodes, edges)
    _ = codeflash_output  # 1.54μs -> 1.38μs (12.1% faster)


def test_edges_and_nodes_can_be_any_iterable_types():
    # The function operates on iterables, not necessarily lists. Use tuples.
    nodes = ({"id": "sink"}, {"id": "other"})
    # create an edge that points from 'other', making 'sink' the last node
    edges = ({"source": "other"},)
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.38μs -> 1.46μs (5.69% slower)


def test_returns_none_when_every_node_has_outgoing_edge():
    # All nodes have outgoing edges -> no last node should be found
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1}, {"source": 2}]
    codeflash_output = find_last_node(nodes, edges)  # 1.54μs -> 1.33μs (15.7% faster)


def test_large_scale_tail_node_identified_correctly():
    # Large scale test with many nodes but kept below 1000 (here N = 700)
    N = 700  # within bounds specified (less than 1000)
    # create nodes with ids 0..N-1 as dicts
    nodes = tuple({"id": i} for i in range(N))
    # create edges that make a long chain: 0 -> 1, 1 -> 2, ..., N-2 -> N-1
    # so the last node should be id N-1
    edges = tuple({"source": i} for i in range(N - 1))
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 8.77ms -> 35.0μs (24935% 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_basic_single_node_no_edges():
    """Test with a single node and no edges - should return that node."""
    nodes = [{"id": "node1", "label": "Start"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.12μs -> 1.00μs (12.5% faster)


def test_basic_linear_chain_three_nodes():
    """Test with a linear chain: node1 -> node2 -> node3."""
    nodes = [
        {"id": "node1", "label": "First"},
        {"id": "node2", "label": "Middle"},
        {"id": "node3", "label": "Last"},
    ]
    edges = [
        {"source": "node1", "target": "node2"},
        {"source": "node2", "target": "node3"},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.38μs (36.4% faster)


def test_basic_two_nodes_one_edge():
    """Test with two nodes where first is source."""
    nodes = [{"id": "node1", "label": "Start"}, {"id": "node2", "label": "End"}]
    edges = [{"source": "node1", "target": "node2"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 1.33μs (12.5% faster)


def test_basic_fan_out_structure():
    """Test with one node having multiple outgoing edges."""
    nodes = [
        {"id": "node1", "label": "Start"},
        {"id": "node2", "label": "Branch1"},
        {"id": "node3", "label": "Branch2"},
    ]
    edges = [
        {"source": "node1", "target": "node2"},
        {"source": "node1", "target": "node3"},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.58μs -> 1.38μs (15.2% faster)


def test_basic_node_with_extra_attributes():
    """Test that nodes with extra attributes are handled correctly."""
    nodes = [
        {"id": "node1", "label": "Start", "type": "input", "color": "red"},
        {"id": "node2", "label": "End", "type": "output", "color": "green"},
    ]
    edges = [{"source": "node1", "target": "node2"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 1.29μs (16.2% faster)


def test_basic_edge_with_extra_attributes():
    """Test that edges with extra attributes don't affect logic."""
    nodes = [{"id": "node1", "label": "Start"}, {"id": "node2", "label": "End"}]
    edges = [
        {"source": "node1", "target": "node2", "weight": 10, "label": "connection"}
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 1.25μs (20.0% faster)


def test_edge_empty_nodes_list():
    """Test with empty nodes list - should return None."""
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 666ns -> 1.00μs (33.4% slower)


def test_edge_empty_edges_list():
    """Test with no edges - first node with no outgoing edges wins."""
    nodes = [{"id": "node1", "label": "First"}, {"id": "node2", "label": "Second"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.08μs -> 959ns (12.9% faster)


def test_edge_all_nodes_are_sources():
    """Test where every node is a source - no last node exists."""
    nodes = [
        {"id": "node1", "label": "Node1"},
        {"id": "node2", "label": "Node2"},
        {"id": "node3", "label": "Node3"},
    ]
    edges = [
        {"source": "node1", "target": "node2"},
        {"source": "node2", "target": "node3"},
        {"source": "node3", "target": "node1"},  # cycle
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.92μs -> 1.46μs (31.4% faster)


def test_edge_cyclic_graph_simple():
    """Test with a simple cycle: node1 -> node2 -> node1."""
    nodes = [{"id": "node1", "label": "A"}, {"id": "node2", "label": "B"}]
    edges = [
        {"source": "node1", "target": "node2"},
        {"source": "node2", "target": "node1"},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.54μs -> 1.38μs (12.1% faster)


def test_edge_disconnected_component():
    """Test with multiple disconnected components."""
    nodes = [
        {"id": "node1", "label": "Component1Start"},
        {"id": "node2", "label": "Component1End"},
        {"id": "node3", "label": "Component2Start"},
        {"id": "node4", "label": "Component2End"},
    ]
    edges = [
        {"source": "node1", "target": "node2"},
        {"source": "node3", "target": "node4"},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.62μs -> 1.33μs (21.9% faster)


def test_edge_node_ids_with_special_characters():
    """Test node IDs with special characters."""
    nodes = [
        {"id": "node-1_v2.0", "label": "Start"},
        {"id": "node/2:special", "label": "End"},
    ]
    edges = [{"source": "node-1_v2.0", "target": "node/2:special"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.67μs -> 1.33μs (25.0% faster)


def test_edge_numeric_node_ids():
    """Test with numeric node IDs (as strings)."""
    nodes = [
        {"id": "1", "label": "First"},
        {"id": "2", "label": "Second"},
        {"id": "3", "label": "Third"},
    ]
    edges = [{"source": "1", "target": "2"}, {"source": "2", "target": "3"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.38μs (36.4% faster)


def test_edge_single_node_in_multiple_edges_as_target():
    """Test where one node is target of multiple edges."""
    nodes = [
        {"id": "node1", "label": "Source1"},
        {"id": "node2", "label": "Source2"},
        {"id": "node3", "label": "Target"},
    ]
    edges = [
        {"source": "node1", "target": "node3"},
        {"source": "node2", "target": "node3"},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.29μs (45.1% faster)


def test_edge_diamond_graph():
    """Test with diamond-shaped graph: node1 -> node2,node3 -> node4."""
    nodes = [
        {"id": "node1", "label": "Top"},
        {"id": "node2", "label": "LeftBranch"},
        {"id": "node3", "label": "RightBranch"},
        {"id": "node4", "label": "Bottom"},
    ]
    edges = [
        {"source": "node1", "target": "node2"},
        {"source": "node1", "target": "node3"},
        {"source": "node2", "target": "node4"},
        {"source": "node3", "target": "node4"},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.50μs -> 1.54μs (62.1% faster)


def test_edge_self_loop():
    """Test where a node has a self-loop."""
    nodes = [{"id": "node1", "label": "Normal"}, {"id": "node2", "label": "SelfLoop"}]
    edges = [
        {"source": "node1", "target": "node2"},
        {"source": "node2", "target": "node2"},  # self-loop
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.67μs -> 1.29μs (28.9% faster)


def test_edge_edge_with_missing_source_field():
    """Test behavior when edge is missing 'source' field - should handle gracefully."""
    nodes = [{"id": "node1", "label": "Start"}, {"id": "node2", "label": "End"}]
    edges = [{"target": "node2"}]  # missing 'source' field
    # This will raise KeyError when accessing e["source"]
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 1.54μs -> 1.00μs (54.2% faster)


def test_edge_none_as_node_id():
    """Test where node ID is None."""
    nodes = [{"id": None, "label": "NoneId"}, {"id": "node2", "label": "RealId"}]
    edges = [{"source": None, "target": "node2"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.33μs (37.4% faster)


def test_edge_empty_string_as_node_id():
    """Test where node ID is an empty string."""
    nodes = [{"id": "", "label": "EmptyId"}, {"id": "node2", "label": "RealId"}]
    edges = [{"source": "", "target": "node2"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.58μs -> 1.29μs (22.6% faster)


def test_edge_unicode_node_ids():
    """Test with unicode characters in node IDs."""
    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.58μs -> 1.29μs (22.5% faster)


def test_large_scale_linear_chain():
    """Test with a large linear chain of 100 nodes."""
    num_nodes = 100
    nodes = [{"id": f"node{i}", "label": f"Node{i}"} for i in range(num_nodes)]
    edges = [
        {"source": f"node{i}", "target": f"node{i+1}"} for i in range(num_nodes - 1)
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 218μs -> 12.6μs (1633% faster)


def test_large_scale_complete_fan_out():
    """Test with one source node connected to 50 sink nodes."""
    source_node = {"id": "source", "label": "Source"}
    sink_nodes = [{"id": f"sink{i}", "label": f"Sink{i}"} for i in range(50)]
    nodes = [source_node] + sink_nodes
    edges = [{"source": "source", "target": f"sink{i}"} for i in range(50)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 3.54μs -> 2.33μs (51.8% faster)


def test_large_scale_complete_fan_in():
    """Test with 50 source nodes all connected to one sink node."""
    source_nodes = [{"id": f"source{i}", "label": f"Source{i}"} for i in range(50)]
    sink_node = {"id": "sink", "label": "Sink"}
    nodes = source_nodes + [sink_node]
    edges = [{"source": f"source{i}", "target": "sink"} for i in range(50)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 62.4μs -> 7.58μs (723% faster)


def test_large_scale_binary_tree():
    """Test with a binary tree structure (63 nodes for depth 6)."""
    # Create binary tree nodes
    nodes = []
    edges = []
    node_count = 0

    # Build tree with depth limit to stay under 1000 nodes
    depth = 6
    for level in range(depth):
        level_start = (2**level) - 1
        level_end = (2 ** (level + 1)) - 1
        for idx in range(level_start, level_end):
            nodes.append({"id": f"node{idx}", "label": f"Node{idx}"})
            if idx > 0:  # Not the root
                parent = (idx - 1) // 2
                edges.append({"source": f"node{parent}", "target": f"node{idx}"})

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 48.0μs -> 6.92μs (593% faster)
    # Check that result is one of the leaf nodes
    leaf_start = (2**depth) - 1


def test_large_scale_multi_path_dag():
    """Test with a DAG having multiple paths (100 nodes, 200 edges)."""
    # Create a layered DAG
    layers = 10
    nodes_per_layer = 10
    nodes = []
    edges = []

    # Create layers
    for layer in range(layers):
        for idx in range(nodes_per_layer):
            node_id = f"layer{layer}_node{idx}"
            nodes.append({"id": node_id, "label": node_id})

    # Connect each node to 2 nodes in next layer
    for layer in range(layers - 1):
        for idx in range(nodes_per_layer):
            source = f"layer{layer}_node{idx}"
            # Connect to nodes in next layer
            for next_idx in [(idx) % nodes_per_layer, (idx + 1) % nodes_per_layer]:
                target = f"layer{layer + 1}_node{next_idx}"
                edges.append({"source": source, "target": target})

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 341μs -> 15.0μs (2186% faster)


def test_large_scale_sparse_graph():
    """Test with a sparse graph (100 nodes, only 10 edges)."""
    nodes = [{"id": f"node{i}", "label": f"Node{i}"} for i in range(100)]
    edges = [{"source": f"node{i}", "target": f"node{i+50}"} for i in range(10)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 5.92μs -> 2.50μs (137% faster)


def test_large_scale_dense_graph():
    """Test with a dense graph where each node connects to the next (99 edges)."""
    num_nodes = 50
    nodes = [{"id": f"node{i}", "label": f"Node{i}"} for i in range(num_nodes)]
    # Each node i connects to node i+1, creating a chain
    edges = [
        {"source": f"node{i}", "target": f"node{i+1}"} for i in range(num_nodes - 1)
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 62.0μs -> 6.54μs (848% faster)


def test_large_scale_all_pairs_except_last():
    """Test with nearly complete graph (all nodes are sources except last)."""
    num_nodes = 20
    nodes = [{"id": f"node{i}", "label": f"Node{i}"} for i in range(num_nodes)]
    edges = []

    # Connect every node to every other node
    for i in range(num_nodes - 1):
        for j in range(i + 1, num_nodes):
            edges.append({"source": f"node{i}", "target": f"node{j}"})

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 97.2μs -> 11.0μs (787% faster)


def test_large_scale_wide_graph():
    """Test with a very wide graph (1 layer source, 1 layer of 100 nodes)."""
    nodes = [{"id": "source", "label": "Source"}]
    for i in range(100):
        nodes.append({"id": f"sink{i}", "label": f"Sink{i}"})

    edges = [{"source": "source", "target": f"sink{i}"} for i in range(100)]

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 5.46μs -> 3.25μs (67.9% faster)


def test_large_scale_deep_graph():
    """Test with a very deep graph (100 nodes in a chain)."""
    nodes = [{"id": f"node{i}", "label": f"Node{i}"} for i in range(100)]
    edges = [{"source": f"node{i}", "target": f"node{i+1}"} for i in range(99)]

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 219μs -> 11.8μs (1763% faster)


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

📊 Performance Profile

View detailed line-by-line performance analysis
To edit these changes git checkout codeflash/optimize-find_last_node-ml2f84v0 and push.