⚡️ Speed up function find_last_node by 24,986%

[codeflash-ai]

📄 24,986% (249.86x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 54.2 milliseconds → 216 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 250x speedup by eliminating a quadratic O(n*m) algorithm and replacing it with a linear O(n+m) approach using set-based lookup.

Key Optimization:

The original code uses a nested loop structure: for each node, it checks ALL edges to verify none have that node as a source. This creates O(n*m) time complexity where n = number of nodes and m = number of edges. The all() call with the inner generator over edges means for each node examined, every edge is checked.

The optimized version:

1. Pre-computes a set of all source node IDs from edges (single O(m) pass)
2. Uses set membership testing (n["id"] not in source_ids) which is O(1) average case, instead of the O(m) linear scan through all edges

Why This is Faster:

• Set lookups in Python are O(1) on average vs O(m) for scanning all edges
• Dictionary/set construction is highly optimized in CPython
• Single pass over edges instead of repeatedly iterating for each node candidate

Impact Based on Test Results:

The speedup scales dramatically with input size:

• Small graphs (2-3 nodes): ~50-70% faster (1-2μs saved)
• Large linear chains (1000 nodes): 323x faster (18ms → 55μs)
• Large cycles (1000 nodes): 328x faster (18.2ms → 55.2μs)

The optimization is most effective when:

• There are many edges to check against each node
• The graph has many nodes but relatively few are "last nodes" (nodes with no outgoing edges)
• The function is called repeatedly in data processing pipelines

This is a classic example of trading space (storing source_ids set) for time, converting repeated linear scans into constant-time lookups.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 20 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 -> 917ns (31.7% faster)


def test_two_nodes_one_edge():
    # Two nodes, one edge from node 1 to node 2: last node should be node 2
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)  # 1.83μs -> 1.17μs (57.2% faster)


def test_three_nodes_linear_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.17μs -> 1.29μs (67.9% faster)


def test_three_nodes_branch():
    # 1 -> 2, 1 -> 3, both 2 and 3 are leaves, return the first found (2)
    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.79μs -> 1.21μs (48.3% faster)


def test_disconnected_nodes():
    # 1 -> 2, 3 is disconnected, should return 2 or 3 (both have no outgoing edges)
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    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():
    # No nodes at all
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 750ns -> 833ns (9.96% slower)


def test_empty_edges():
    # Multiple nodes, no edges: first node should be returned
    nodes = [{"id": "a"}, {"id": "b"}, {"id": "c"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 1.25μs -> 958ns (30.5% faster)


def test_all_nodes_with_outgoing_edges():
    # Every node has an outgoing edge, so no last node
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)  # 1.88μs -> 1.21μs (55.2% faster)


def test_cycle():
    # 1 -> 2 -> 3 -> 1 (cycle), no last node
    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)  # 2.25μs -> 1.29μs (74.1% faster)


def test_duplicate_edges():
    # 1 -> 2 (twice), 2 is still the last node
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)  # 1.88μs -> 1.21μs (55.2% faster)


def test_self_loop():
    # 1 -> 1 (self loop), should not be last node
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)  # 1.75μs -> 1.12μs (55.6% faster)


def test_node_with_multiple_outgoing_edges():
    # 1 -> 2, 1 -> 3, 2 and 3 are leaves
    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.83μs -> 1.17μs (57.1% faster)


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


def test_node_ids_are_mixed_types():
    # Node IDs are mixed types (should work as long as they match)
    nodes = [{"id": 1}, {"id": "2"}]
    edges = [{"source": 1, "target": "2"}]
    codeflash_output = find_last_node(nodes, edges)  # 1.88μs -> 1.08μs (73.1% faster)


def test_edge_with_nonexistent_node():
    # Edge references a node not in nodes (should ignore)
    nodes = [{"id": 1}]
    edges = [{"source": 1, "target": 2}]  # node 2 does not exist
    codeflash_output = find_last_node(nodes, edges)  # 1.38μs -> 1.08μs (27.0% faster)


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


def test_large_linear_chain():
    # Large linear chain: 0 -> 1 -> 2 -> ... -> 999, last node is 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.0ms -> 55.6μs (32325% faster)


def test_large_branching_graph():
    # Star graph: 0 -> 1, 0 -> 2, ..., 0 -> 999, all leaves are 1..999
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": 0, "target": i} for i in range(1, N)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 37.7μs -> 20.0μs (88.9% faster)


def test_large_disconnected_graph():
    # 500 isolated nodes, 500 in a chain: 0..499 (isolated), 500->501->...->999
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i + 1} for i in range(500, N - 1)]
    # The isolated nodes are also leaves, so any of them or the last in chain is valid
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 20.4μs -> 13.6μs (49.5% faster)
    valid_leaves = [{"id": i} for i in range(500)] + [{"id": N - 1}]


def test_large_cycle():
    # Large cycle: 0->1->2->...->999->0, no last node
    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)  # 18.2ms -> 55.2μs (32842% faster)


def test_large_graph_with_self_loops():
    # Each node has a self loop, no last node
    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)  # 17.9ms -> 54.7μs (32548% 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-mjj06nsd and push.