⚡️ Speed up function find_last_node by 4,605%

[codeflash-ai]

📄 4,605% (46.05x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 2.98 milliseconds → 63.4 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 46x speedup (4605% improvement) by eliminating the quadratic O(N×M) complexity of the original implementation and replacing it with a linear O(N+M) approach.

Key Optimization:

The original code uses nested iteration via all(e["source"] != n["id"] for e in edges), which re-scans the entire edges collection for every node. This creates O(N×M) time complexity where N is the number of nodes and M is the number of edges.

The optimized version introduces a smart branching strategy:

1. For re-iterable collections (lists, tuples - the common case): It builds a set of source IDs in one pass (sources = {e["source"] for e in edges}), then checks membership with O(1) lookups. This reduces complexity to O(N+M) with minimal O(M) memory overhead.

2. For single-pass iterators: It preserves the original consumption semantics by advancing the iterator sequentially across nodes, avoiding re-iteration while maintaining correctness for iterator inputs.

Why This Works:

• Set membership testing (n["id"] not in sources) is O(1) average case, versus O(M) for the all() scan
• Building the set requires one pass through edges (O(M)), then checking all nodes (O(N))
• Total: O(N+M) vs O(N×M)

Impact on Test Cases:

The speedup scales dramatically with graph size:

• Small graphs (2-3 nodes): 45-109% faster
• Medium chains (100 nodes): 3026% faster
• Large trees (511 nodes): 13121% faster
• Dense graphs (20 fully-connected nodes): 1403% faster
• Complex patterns (100-node diamonds): 3495% faster

The optimization is particularly effective for larger graphs or graphs with many edges, where the quadratic behavior of the original becomes prohibitive. Even small graphs benefit from the reduced overhead of set lookups versus repeated linear scans.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

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 - node should be returned as last node."""
    nodes = [{"id": 1, "name": "Node1"}]
    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: Node1 -> Node2 -> Node3. Node3 should be returned."""
    nodes = [
        {"id": 1, "name": "Node1"},
        {"id": 2, "name": "Node2"},
        {"id": 3, "name": "Node3"},
    ]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 958ns (91.3% faster)


def test_find_last_node_two_nodes_one_edge():
    """Test with two nodes and one edge connecting them."""
    nodes = [{"id": 1, "name": "Start"}, {"id": 2, "name": "End"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 833ns (80.1% faster)


def test_find_last_node_multiple_start_nodes():
    """Test with multiple nodes without outgoing edges (should return first one)."""
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}, {"id": 3, "name": "C"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.42μs -> 791ns (79.1% faster)


def test_find_last_node_with_additional_node_attributes():
    """Test that extra node attributes are preserved in the returned node."""
    nodes = [
        {"id": 1, "name": "Start", "type": "input", "value": 42},
        {"id": 2, "name": "End", "type": "output", "value": 100},
    ]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 833ns (75.0% faster)


def test_find_last_node_branching_flow():
    """Test with a branching flow where one branch is longer."""
    nodes = [
        {"id": 1, "name": "Start"},
        {"id": 2, "name": "Branch1"},
        {"id": 3, "name": "Branch2"},
        {"id": 4, "name": "End1"},
    ]
    edges = [
        {"source": 1, "target": 2},
        {"source": 1, "target": 3},
        {"source": 2, "target": 4},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 959ns (95.5% faster)


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


def test_find_last_node_empty_edges_list():
    """Test with empty edges list - first node should be returned as last node."""
    nodes = [{"id": 1, "name": "OnlyNode"}]
    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_have_outgoing_edges():
    """Test with a cyclic graph where all nodes have outgoing edges."""
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}, {"id": 3, "name": "C"}]
    edges = [
        {"source": 1, "target": 2},
        {"source": 2, "target": 3},
        {"source": 3, "target": 1},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 958ns (91.3% faster)


def test_find_last_node_with_self_loop():
    """Test with a node that has a self-loop edge."""
    nodes = [{"id": 1, "name": "NodeWithLoop"}, {"id": 2, "name": "NormalNode"}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.42μs -> 833ns (70.1% faster)


def test_find_last_node_with_string_ids():
    """Test with string node IDs instead of integers."""
    nodes = [
        {"id": "nodeA", "name": "A"},
        {"id": "nodeB", "name": "B"},
        {"id": "nodeC", "name": "C"},
    ]
    edges = [
        {"source": "nodeA", "target": "nodeB"},
        {"source": "nodeB", "target": "nodeC"},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.92μs -> 917ns (109% faster)


def test_find_last_node_with_duplicate_edges():
    """Test with duplicate edges between same nodes."""
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 917ns (63.6% faster)


def test_find_last_node_with_extra_edge_attributes():
    """Test that extra edge attributes don't affect the result."""
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}]
    edges = [{"source": 1, "target": 2, "weight": 5, "label": "connection"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 833ns (75.0% faster)


def test_find_last_node_none_as_node_id():
    """Test with None as a node ID value."""
    nodes = [{"id": None, "name": "NullNode"}, {"id": 1, "name": "RealNode"}]
    edges = [{"source": 1, "target": None}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 709ns (76.3% faster)


def test_find_last_node_node_with_zero_id():
    """Test with zero as a node ID."""
    nodes = [{"id": 0, "name": "ZeroNode"}, {"id": 1, "name": "OneNode"}]
    edges = [{"source": 0, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 834ns (74.8% faster)


def test_find_last_node_negative_ids():
    """Test with negative node IDs."""
    nodes = [{"id": -1, "name": "NegOne"}, {"id": -2, "name": "NegTwo"}]
    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_with_float_ids():
    """Test with float node IDs."""
    nodes = [{"id": 1.5, "name": "Float1"}, {"id": 2.5, "name": "Float2"}]
    edges = [{"source": 1.5, "target": 2.5}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.58μs -> 917ns (72.6% faster)


def test_find_last_node_returns_none_when_all_nodes_source():
    """Test that None is returned when all nodes are sources (cyclic or all have outgoing)."""
    nodes = [{"id": "A", "name": "A"}, {"id": "B", "name": "B"}]
    edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "A"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.54μs -> 875ns (76.2% faster)


def test_find_last_node_large_linear_chain():
    """Test with a large linear chain of 100 nodes."""
    # Create a chain: 1 -> 2 -> 3 -> ... -> 100
    nodes = [{"id": i, "name": f"Node{i}"} for i in range(1, 101)]
    edges = [{"source": i, "target": i + 1} for i in range(1, 100)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 201μs -> 6.46μs (3026% faster)


def test_find_last_node_large_branching_tree():
    """Test with a large branching tree structure (511 nodes, 2-level binary tree)."""
    # Create a tree where node 1 branches to 2,3; 2 branches to 4,5; etc.
    nodes = [{"id": i, "name": f"Node{i}"} for i in range(1, 512)]
    edges = []
    for i in range(1, 256):
        edges.append({"source": i, "target": 2 * i})
        edges.append({"source": i, "target": 2 * i + 1})

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.30ms -> 17.4μs (13121% faster)


def test_find_last_node_completely_connected_graph():
    """Test with a fully connected graph (every node connects to every other node)."""
    # Create 20 nodes with edges from each to all others
    nodes = [{"id": i, "name": f"N{i}"} for i in range(20)]
    edges = []
    for i in range(20):
        for j in range(20):
            if i != j:
                edges.append({"source": i, "target": j})

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 132μs -> 8.83μs (1403% faster)


def test_find_last_node_star_topology():
    """Test with a star topology (central node connected to many peripheral nodes)."""
    # Node 0 is center, connects to 1-49 (50 peripheral nodes)
    nodes = [{"id": i, "name": f"Node{i}"} for i in range(50)]
    edges = [{"source": 0, "target": i} for i in range(1, 50)]

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 3.38μs -> 1.75μs (92.9% faster)


def test_find_last_node_multiple_disconnected_chains():
    """Test with multiple separate chains (disconnected graph)."""
    # Chain 1: 1 -> 2 -> 3
    # Chain 2: 4 -> 5 -> 6
    # Chain 3: 7 -> 8
    nodes = [{"id": i, "name": f"N{i}"} for i in range(1, 9)]
    edges = [
        {"source": 1, "target": 2},
        {"source": 2, "target": 3},
        {"source": 4, "target": 5},
        {"source": 5, "target": 6},
        {"source": 7, "target": 8},
    ]

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 1.08μs (80.8% faster)


def test_find_last_node_with_many_isolated_nodes():
    """Test with many isolated nodes (no edges)."""
    # Create 200 nodes with no edges
    nodes = [{"id": i, "name": f"Isolated{i}"} for i in range(200)]
    edges = []

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.04μs -> 667ns (56.1% faster)


def test_find_last_node_long_chain_with_cycles():
    """Test with a long chain plus a cycle back to beginning."""
    # 1 -> 2 -> 3 -> ... -> 50 -> 1 (cycle back)
    nodes = [{"id": i, "name": f"N{i}"} for i in range(1, 51)]
    edges = [{"source": i, "target": i + 1} for i in range(1, 50)]
    edges.append({"source": 50, "target": 1})  # cycle back

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


def test_find_last_node_diamond_pattern_large():
    """Test with multiple diamond patterns (node 1 -> 2,3 -> 4; 4 -> 5,6 -> 7; etc.)."""
    # Create cascading diamonds: total 100 nodes
    nodes = [{"id": i, "name": f"N{i}"} for i in range(1, 101)]
    edges = []

    # Build diamond patterns: 1->[2,3]->4, 4->[5,6]->7, etc.
    current_top = 1
    nodes_used = 4
    while nodes_used <= 100:
        edges.append({"source": current_top, "target": current_top + 1})
        edges.append({"source": current_top, "target": current_top + 2})
        edges.append({"source": current_top + 1, "target": current_top + 3})
        edges.append({"source": current_top + 2, "target": current_top + 3})
        current_top = current_top + 3
        nodes_used += 3

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 259μs -> 7.21μs (3495% faster)


def test_find_last_node_sparse_large_graph():
    """Test with a sparse large graph (1000 nodes, few edges)."""
    # Create 1000 nodes
    nodes = [{"id": i, "name": f"Node{i}"} for i in range(1000)]

    # Create sparse edges: only connect every 100th node
    edges = []
    for i in range(0, 900, 100):
        edges.append({"source": i, "target": i + 1})

    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.08μs -> 1.25μs (66.6% 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-mlf9sxm8 and push.