⚡️ Speed up function find_last_node by 21,185%

[codeflash-ai]

📄 21,185% (211.85x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 93.7 milliseconds → 440 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 21,000% speedup by eliminating a severe O(n×m) algorithmic bottleneck through pre-computing a set of source node IDs.

Key Optimization:

The original implementation used a nested iteration pattern:

all(e["source"] != n["id"] for e in edges)

For each node, it checked all edges to verify the node wasn't a source. With n nodes and m edges, this resulted in O(n×m) comparisons—catastrophic for larger graphs.

The optimized version pre-computes a hash set of source IDs once:

sources = {e["source"] for e in edges}
return next((n for n in nodes if n["id"] not in sources), None)

This reduces the algorithm to O(m + n):

• O(m) to build the sources set (one pass through edges)
• O(n) for the lookup loop, with O(1) average-case hash set membership checks

Why This Works:

Python's set implementation uses hash tables, providing constant-time lookups versus the linear scan required by all(). The line profiler shows the dramatic impact:

• Original: 697ms total, all spent in the nested iteration
• Optimized: 707μs total (321μs building set, 386μs finding node)

Performance Characteristics:

The optimization excels particularly on large graphs:

• Large linear chain (1000 nodes): 19.5ms → 98μs (19,822% faster)
• Large cycle graph (1000 nodes): 19.7ms → 98.9μs (19,851% faster)

Small graphs show modest improvements (30-95% faster) since overhead is dominated by Python's interpreter rather than the algorithm. The only slight regression is empty inputs (9-23% slower) where set creation overhead isn't amortized, but this is negligible at sub-microsecond scales.

Impact: If find_last_node is called in graph processing pipelines or hot paths, this optimization will dramatically reduce execution time, especially for graphs with hundreds or thousands of nodes/edges.

✅ 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

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

# unit tests

# 1. Basic Test Cases


def test_single_node_no_edges():
    # Only one node, no edges; should return the node itself
    nodes = [{"id": "A"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 958ns (30.5% faster)


def test_two_nodes_one_edge():
    # Two nodes, one edge from A -> B; B is the last node
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 1.12μs (74.0% faster)


def test_three_nodes_linear_chain():
    # Three nodes in a chain: A -> B -> C; C is the last node
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "C"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.33μs -> 1.21μs (93.1% faster)


def test_multiple_last_nodes_returns_first():
    # Two nodes with no outgoing edges (B, C); should return B (first in list)
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.79μs -> 1.12μs (59.3% faster)


# 2. 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 -> 875ns (9.60% slower)


def test_nodes_no_edges():
    # Multiple nodes, no edges; should return the first node
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 1.00μs (25.0% faster)


def test_all_nodes_have_outgoing_edges():
    # Each node has an outgoing edge; no last node, should return None
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [
        {"source": "A", "target": "B"},
        {"source": "B", "target": "C"},
        {"source": "C", "target": "A"},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.33μs -> 1.29μs (80.7% faster)


def test_edge_points_to_nonexistent_node():
    # Edge points to a node not in nodes; should not affect result
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [
        {"source": "A", "target": "B"},
        {"source": "B", "target": "C"},
    ]  # C not in nodes
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.21μs (55.2% faster)


def test_node_with_self_loop():
    # Node with a self-loop should not be a last node
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "A"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.12μs (66.7% faster)


def test_duplicate_node_ids():
    # Duplicate node ids; should return the first with no outgoing edge
    nodes = [{"id": "A"}, {"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.12μs -> 1.17μs (82.1% faster)


def test_edge_with_missing_source_key():
    # Edge missing 'source' key; should not throw error, just ignore
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"target": "B"}]
    with pytest.raises(KeyError):
        # The implementation expects 'source' to be present
        find_last_node(nodes, edges)  # 1.54μs -> 792ns (94.7% faster)


def test_node_with_extra_keys():
    # Nodes with extra data; should still return correct node
    nodes = [{"id": "A", "label": "Start"}, {"id": "B", "label": "End"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.12μs (63.0% faster)


def test_edge_with_extra_keys():
    # Edges with extra keys should not affect result
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B", "weight": 5}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.79μs -> 1.12μs (59.3% faster)


def test_nodes_with_non_string_ids():
    # Node ids are integers
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.17μs (60.7% faster)


def test_nodes_with_mixed_type_ids():
    # Node ids are mixed types, should match by equality
    nodes = [{"id": 1}, {"id": "2"}]
    edges = [{"source": 1, "target": "2"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.12μs (66.7% faster)


def test_nodes_with_none_id():
    # Node with id None; should be handled correctly
    nodes = [{"id": None}, {"id": "A"}]
    edges = [{"source": "A", "target": None}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.54μs -> 1.08μs (42.3% faster)


def test_edges_with_none_source():
    # Edge with source None; should not match any node except node with id None
    nodes = [{"id": "A"}, {"id": None}]
    edges = [{"source": None, "target": "A"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.54μs -> 1.04μs (48.0% faster)


# 3. Large Scale Test Cases


def test_large_linear_chain():
    # 1000 nodes in a linear chain; last node should be the last in the list
    N = 1000
    nodes = [{"id": str(i)} for i in range(N)]
    edges = [{"source": str(i), "target": str(i + 1)} for i in range(N - 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 19.5ms -> 98.0μs (19822% faster)


def test_large_star_graph():
    # One central node with edges to 999 leaves; all leaves are last nodes, should return first leaf
    N = 1000
    nodes = [{"id": "center"}] + [{"id": f"leaf_{i}"} for i in range(N - 1)]
    edges = [{"source": "center", "target": f"leaf_{i}"} for i in range(N - 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 40.8μs -> 19.2μs (112% faster)


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


def test_large_all_nodes_have_outgoing_edges():
    # All nodes have outgoing edges, forming a cycle; should return None
    N = 1000
    nodes = [{"id": str(i)} for i in range(N)]
    edges = [{"source": str(i), "target": str((i + 1) % N)} for i in range(N)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 19.7ms -> 98.9μs (19851% 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
from src.algorithms.graph import find_last_node

# unit tests

# --- Basic Test Cases ---


def test_single_node_no_edges():
    # A single node and no edges: should return the node itself
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.29μs -> 958ns (34.8% faster)


def test_two_nodes_one_edge():
    # Two nodes, one edge from 1 -> 2: should return node 2 (no outgoing edges)
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 1.17μs (67.8% faster)


def test_three_nodes_linear_chain():
    # 1 -> 2 -> 3, should return node 3
    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.33μs -> 1.21μs (93.0% faster)


def test_three_nodes_multiple_leaves():
    # 1 -> 2, 1 -> 3, should return node 2 (first leaf found)
    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.21μs (58.7% faster)


def test_three_nodes_disconnected():
    # 1 -> 2, node 3 is disconnected (should be returned)
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.12μs (62.9% faster)
    # If we reorder nodes, node 3 is first leaf
    nodes2 = [{"id": 1}, {"id": 3}, {"id": 2}]
    codeflash_output = find_last_node(nodes2, edges)
    result2 = codeflash_output  # 916ns -> 500ns (83.2% 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  # 708ns -> 916ns (22.7% slower)


def test_nodes_with_self_loop():
    # Node with a self-loop should not be considered a leaf
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.08μs (69.3% faster)


def test_all_nodes_have_outgoing_edges():
    # All nodes have outgoing edges, no leaf: should return None
    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.96μs -> 1.21μs (62.1% faster)


def test_duplicate_edges():
    # Duplicated edges should not affect result
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [
        {"source": 1, "target": 2},
        {"source": 1, "target": 2},
        {"source": 2, "target": 3},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.42μs -> 1.29μs (87.0% faster)


def test_multiple_leaf_nodes_order():
    # Multiple leaf nodes, should return first in nodes order
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.12μs (66.7% faster)


def test_nodes_with_non_integer_ids():
    # Node ids are strings
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "C"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.38μs -> 1.21μs (96.6% faster)


def test_edges_with_extra_fields():
    # Edges may have extra fields, should be ignored
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2, "weight": 5}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.12μs (62.9% faster)


def test_nodes_with_extra_fields():
    # Nodes may have extra fields, should be returned as-is
    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.83μs -> 1.12μs (62.9% faster)


def test_edge_with_missing_source_field():
    # Edge missing 'source' field should not break the function
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"target": 2}]
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 1.58μs -> 750ns (111% faster)


def test_edge_with_missing_target_field():
    # Edge missing 'target' field is fine, since we only care about 'source'
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.12μs (62.9% faster)


# --- Large Scale Test Cases ---


def test_large_linear_chain():
    # Large chain: 1 -> 2 -> 3 -> ... -> 999
    N = 999
    nodes = [{"id": i} for i in range(1, N + 1)]
    edges = [{"source": i, "target": i + 1} for i in range(1, N)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.1ms -> 55.9μs (32256% faster)


def test_large_star_graph():
    # Node 0 connects to all others, all others are leaves
    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.2μs -> 20.1μs (85.5% faster)


def test_large_disconnected_nodes():
    # All nodes, no edges: first node is leaf
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.33μs -> 1.04μs (28.0% faster)


def test_large_cycle_graph():
    # Each node i connects to (i+1)%N, so all nodes have outgoing edges, no leaves
    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)
    result = codeflash_output  # 18.2ms -> 54.8μs (33181% faster)


def test_large_graph_with_multiple_leaves():
    # Nodes 0..997 have outgoing edges, 998 and 999 are leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i + 1} for i in range(N - 2)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.0ms -> 55.7μs (32248% 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-mjmt8n0l and push.