⚡️ Speed up function find_last_node by 19,149%

[codeflash-ai]

📄 19,149% (191.49x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 93.6 milliseconds → 486 microseconds (best of 178 runs)

📝 Explanation and details

The optimized code achieves a 191x speedup by eliminating a nested loop antipattern that creates O(n*m) time complexity, where n is the number of nodes and m is the number of edges.

Key Optimization:

The original code uses a nested generator expression: for each node, it checks all(e["source"] != n["id"] for e in edges), which means it iterates through ALL edges for EVERY node. This creates O(n*m) comparisons.

The optimized version preprocesses edges into a set (source_ids = set(e["source"] for e in edges)) once, then performs O(1) membership checks for each node. This reduces complexity to O(m + n).

Why This Works:

1. Set lookup is O(1) vs iterating through all edges which is O(m)
2. Single pass through edges instead of m passes (once per node)
3. Changed to n.get("id") instead of n["id"] to handle nodes missing the 'id' field gracefully, preserving the original's behavior of returning nodes without accessing missing keys

Performance Impact by Test Case:

• Small graphs (2-3 nodes): 20-55% faster - modest gains since overhead is minimal
• Large chains/cycles (1000 nodes): 24,000%+ faster - dramatic speedup where nested iteration becomes prohibitively expensive
• Sparse graphs: 172% faster - benefits scale with edge density
• Star graphs: 25% faster - edges processed once instead of repeatedly

Best Use Cases:

The optimization excels when:

• Large number of nodes AND edges (test_large_linear_chain: 24,406% faster)
• Dense edge connectivity (test_large_all_connected: 24,432% faster)
• Called frequently in graph traversal algorithms where this function might be on a hot path

The optimization maintains correctness while transforming a quadratic algorithm into linear time.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 38 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)  # 1.21μs -> 1.21μs (0.000% faster)


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


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


def test_multiple_end_nodes():
    # Two end nodes (no outgoing edges): should return the first one found
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}]
    # Both B and C have no outgoing edges; function returns B (first in nodes)
    codeflash_output = find_last_node(nodes, edges)  # 1.96μs -> 1.46μs (34.3% faster)


def test_node_with_self_loop():
    # Node with a self-loop is not a last node
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 1}]
    # Node 2 has no outgoing edges, should be returned
    codeflash_output = find_last_node(nodes, edges)  # 1.79μs -> 1.38μs (30.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)  # 791ns -> 1.00μs (20.9% slower)


def test_edges_but_no_nodes():
    # Edges but no nodes: should return None
    nodes = []
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)  # 750ns -> 1.21μs (37.9% slower)


def test_nodes_with_no_matching_edges():
    # Nodes with edges referencing IDs not in nodes
    nodes = [{"id": 10}, {"id": 20}]
    edges = [{"source": 30, "target": 40}]
    # Both nodes have no outgoing edges
    codeflash_output = find_last_node(nodes, edges)  # 1.46μs -> 1.38μs (6.11% faster)


def test_all_nodes_have_outgoing_edges():
    # All nodes have outgoing edges: should return None
    nodes = [{"id": "x"}, {"id": "y"}]
    edges = [{"source": "x", "target": "y"}, {"source": "y", "target": "x"}]
    codeflash_output = find_last_node(nodes, edges)  # 2.00μs -> 1.58μs (26.3% faster)


def test_duplicate_node_ids():
    # Duplicate node IDs: function should return the first with no outgoing edge
    nodes = [{"id": 1}, {"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    # Node with id 2 has no outgoing edge
    codeflash_output = find_last_node(nodes, edges)  # 2.04μs -> 1.50μs (36.1% faster)


def test_node_id_is_none():
    # Node with id=None, edge with source=None
    nodes = [{"id": None}, {"id": 1}]
    edges = [{"source": None, "target": 1}]
    # Node with id=1 has no outgoing edge
    codeflash_output = find_last_node(nodes, edges)  # 2.00μs -> 1.42μs (41.1% faster)


def test_node_id_with_different_types():
    # Node IDs are of different types (int, str)
    nodes = [{"id": 1}, {"id": "1"}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    # Node "1" and 2 have no outgoing edges; function returns "1" (first in nodes after 1)
    codeflash_output = find_last_node(nodes, edges)  # 2.00μs -> 1.54μs (29.7% faster)


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


def test_large_chain():
    # Large chain of 1000 nodes: 0->1->2->...->999
    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)  # 18.8ms -> 75.5μs (24816% faster)


def test_large_star_graph():
    # Star graph: node 0 points to all others, last nodes are 1..999
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": 0, "target": i} for i in range(1, N)]
    # First node with no outgoing edges is 1
    codeflash_output = find_last_node(nodes, edges)  # 39.4μs -> 31.4μs (25.3% faster)


def test_large_no_edges():
    # 1000 nodes, no edges: should return the first node
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 1.38μs -> 1.29μs (6.42% faster)


def test_large_all_nodes_have_outgoing_edges():
    # 1000 nodes, each node points to itself (self-loop), so all have outgoing edges
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i} for i in range(N)]
    codeflash_output = find_last_node(nodes, edges)  # 18.6ms -> 75.7μs (24447% faster)


def test_large_graph_multiple_last_nodes():
    # 1000 nodes, nodes 0-998 point to 999, so 999 has no outgoing edges
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": N - 1} for i in range(N - 1)]
    codeflash_output = find_last_node(nodes, edges)  # 18.8ms -> 75.5μs (24736% 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

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


def test_single_node_no_edges():
    # Only one node, no edges. That node is the last node.
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 1.25μs (0.000% faster)


def test_two_nodes_one_edge():
    # Node 1 points to Node 2. Node 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.88μs -> 1.46μs (28.6% faster)


def test_three_nodes_chain():
    # 1 -> 2 -> 3. Node 3 is the 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.29μs -> 1.62μs (41.0% faster)


def test_multiple_end_nodes():
    # 1 -> 2, 1 -> 3. Both 2 and 3 are possible last nodes; function returns first it finds.
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.92μs -> 1.58μs (21.1% faster)


def test_all_nodes_have_outgoing_edges():
    # Each node has an outgoing edge, so there is no last node.
    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.88μs -> 1.54μs (21.6% 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  # 791ns -> 1.08μs (27.0% slower)


def test_edges_but_no_nodes():
    # Edges exist but nodes list is empty. Should return None.
    nodes = []
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 750ns -> 1.25μs (40.0% slower)


def test_edges_with_missing_source_field():
    # Edges missing 'source' field should raise KeyError.
    nodes = [{"id": 1}]
    edges = [{"src": 1, "target": 1}]
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 1.88μs -> 1.17μs (60.7% faster)


def test_edges_with_extra_fields():
    # Edges with extra fields should be ignored.
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2, "weight": 10}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.50μs (22.2% faster)


def test_nodes_with_extra_fields():
    # Nodes with extra fields should be handled correctly.
    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.79μs -> 1.46μs (22.8% faster)


def test_duplicate_node_ids():
    # Duplicate node IDs: function should still return a node with no outgoing edge.
    nodes = [{"id": 1}, {"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.08μs -> 1.46μs (42.8% faster)


def test_node_id_not_in_edges():
    # Edges refer to a node id not in nodes list; should not affect result.
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 3, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 1.38μs (6.11% faster)


def test_cycle_graph():
    # 1 -> 2 -> 3 -> 1, cycle. All nodes have outgoing edges. Should return None.
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [
        {"source": 1, "target": 2},
        {"source": 2, "target": 3},
        {"source": 3, "target": 1},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.21μs -> 1.58μs (39.5% faster)


def test_disconnected_graph():
    # 1 -> 2, 3 (no edges). Node 2 and 3 are possible last nodes.
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.71μs -> 1.38μs (24.2% faster)


def test_non_integer_node_ids():
    # Node IDs are strings.
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 1.50μs (30.5% faster)


def test_mixed_type_node_ids():
    # Node IDs are mixed types; should work as long as equality is respected.
    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.58μs -> 1.67μs (55.0% faster)


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


def test_large_linear_chain():
    # Large chain: 1 -> 2 -> ... -> 1000
    nodes = [{"id": i} for i in range(1, 1001)]
    edges = [{"source": i, "target": i + 1} for i in range(1, 1000)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.7ms -> 76.2μs (24406% faster)


def test_large_star_graph():
    # Node 0 points to all others, all others have no outgoing edges.
    nodes = [{"id": 0}] + [{"id": i} for i in range(1, 1001)]
    edges = [{"source": 0, "target": i} for i in range(1, 1001)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 39.4μs -> 31.3μs (25.8% faster)


def test_large_all_connected():
    # Each node points to the next, last node points to first (cycle).
    nodes = [{"id": i} for i in range(1000)]
    edges = [{"source": i, "target": (i + 1) % 1000} for i in range(1000)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.7ms -> 76.0μs (24432% faster)


def test_large_sparse_graph():
    # 1000 nodes, only 10 edges: many last nodes.
    nodes = [{"id": i} for i in range(1000)]
    edges = [{"source": i, "target": i + 1} for i in range(10)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 6.58μs -> 2.42μs (172% faster)


def test_large_disconnected_graph():
    # 1000 nodes, no edges. Any node is a last node.
    nodes = [{"id": i} for i in range(1000)]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.33μs -> 1.33μs (0.000% 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-mjizyrzy and push.