⚡️ Speed up function find_last_node by 18,969%

[codeflash-ai]

📄 18,969% (189.69x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 75.3 milliseconds → 395 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 190x speedup by eliminating a critical algorithmic inefficiency: replacing an O(N×M) nested iteration with O(N+M) operations using a set-based lookup.

Key Optimization

Original approach: For each node, iterate through ALL edges to check if that node is a source

• Time complexity: O(N×M) where N = nodes, M = edges
• The all(e["source"] != n["id"] for e in edges) check runs M comparisons for each of the N nodes

Optimized approach: Pre-build a set of all edge sources once, then check membership

• Time complexity: O(N+M)
• Build edge_sources set in O(M) time
• Check each node's membership in O(1) average time per node

Why This Matters

The performance difference becomes dramatic as graph size increases:

• Small graphs (2-3 nodes): 30-75% faster due to reduced iteration overhead
• Linear chains (1000 nodes): 290x faster (18ms → 62μs) because the original code performs ~500,000 comparisons vs. optimized doing ~2,000 operations
• Fully connected graphs (50 nodes, 2,450 edges): 40x faster (2.14ms → 51.8μs) where the quadratic behavior is most punishing

Additional Change

The code uses n.get("id") instead of n["id"] to handle nodes missing the "id" key gracefully, maintaining the same behavior as the original code which only accessed n["id"] during the comparison check. This prevents KeyError exceptions when processing malformed node data while preserving the performance benefit.

The optimization is universally beneficial across all test cases (except trivially empty graphs where overhead slightly increases by 5-9%). It's especially impactful for graphs with many edges or when checking nodes late in the iteration order.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

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():
    # One node, no edges: should return the node itself
    nodes = [{"id": 1, "label": "A"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.29μs -> 1.00μs (29.2% faster)


def test_two_nodes_one_edge():
    # Two nodes, one edge from first to second: should return the second node
    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.88μs -> 1.17μs (60.7% faster)


def test_three_nodes_linear_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)
    result = codeflash_output  # 2.17μs -> 1.29μs (67.7% faster)


def test_multiple_possible_last_nodes():
    # Two nodes, no edges: both are "last nodes", should return the first in nodes
    nodes = [{"id": 10}, {"id": 20}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 1.00μs (25.0% faster)


def test_cycle_graph():
    # 1->2->3->1 (cycle): No last node, 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.25μs -> 1.33μs (68.8% faster)


# ---------------- EDGE TEST CASES ----------------


def test_empty_nodes_and_edges():
    # Empty graph: should return None
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 792ns -> 875ns (9.49% slower)


def test_edges_with_nonexistent_source():
    # Edges referencing non-existent nodes: should still return node with no outgoing edge
    nodes = [{"id": 1}]
    edges = [{"source": 2, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.54μs -> 1.17μs (32.1% faster)


def test_multiple_last_nodes_returns_first():
    # Multiple nodes with no outgoing edges, returns the first one in nodes order
    nodes = [{"id": "x"}, {"id": "y"}, {"id": "z"}]
    edges = [{"source": "x", "target": "y"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.04μs -> 1.17μs (75.1% faster)


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


def test_edge_with_missing_source_key():
    # Edge missing 'source' key: should raise KeyError
    nodes = [{"id": 1}]
    edges = [{"target": 1}]
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 1.83μs -> 792ns (132% faster)


def test_edge_with_missing_target_key():
    # Edge missing 'target' key: should not affect last node detection
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.17μs (57.1% faster)


def test_non_integer_ids():
    # Node ids are strings
    nodes = [{"id": "nodeA"}, {"id": "nodeB"}]
    edges = [{"source": "nodeA", "target": "nodeB"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 1.17μs (67.8% faster)


# ---------------- LARGE SCALE TEST CASES ----------------


def test_large_linear_chain():
    # 1000 nodes in a linear chain: last node is the last in the list
    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)
    result = codeflash_output  # 18.4ms -> 62.8μs (29266% faster)


def test_large_star_graph():
    # 1 root node with 999 outgoing edges to 999 leaves
    N = 1000
    nodes = [{"id": 0}] + [{"id": i} for i in range(1, N)]
    edges = [{"source": 0, "target": i} for i in range(1, N)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 38.4μs -> 20.5μs (87.6% faster)


def test_large_fully_connected_graph():
    # Every node has an edge to every other node (except self): no last node
    N = 50  # 50*49 = 2450 edges, under 1000 node/edge limit
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": j} for i in range(N) for j in range(N) if i != j]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.14ms -> 51.8μs (4029% faster)


def test_large_disconnected_graph():
    # 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)
    result = codeflash_output  # 1.33μs -> 1.08μs (23.1% faster)


def test_large_graph_with_multiple_last_nodes():
    # 1000 nodes, first 990 in a chain, last 10 isolated
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i + 1} for i in range(989)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.0ms -> 62.1μs (28860% 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():
    # One node, no edges: should return the node itself
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.29μs -> 1.00μs (29.1% 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.79μs -> 1.17μs (53.6% faster)


def test_three_nodes_chain():
    # 1->2->3, should return 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.29μs -> 1.33μs (71.7% faster)


def test_multiple_nodes_multiple_sinks():
    # 1->2, 1->3, 2->4, 3->5 (4 and 5 are both sinks, should return 4 as it's first in nodes)
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}, {"id": 5}]
    edges = [
        {"source": 1, "target": 2},
        {"source": 1, "target": 3},
        {"source": 2, "target": 4},
        {"source": 3, "target": 5},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.71μs -> 1.42μs (91.1% faster)


def test_no_sink_nodes():
    # All nodes have outgoing edges (cycle), 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.25μs -> 1.29μs (74.1% 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 -> 834ns (5.16% slower)


def test_edges_with_missing_nodes():
    # Edges refer to node IDs not present in nodes list
    nodes = [{"id": 1}]
    edges = [{"source": 2, "target": 3}]
    # Node 1 has no outgoing edges, should return node 1
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.50μs -> 1.12μs (33.3% faster)


def test_edges_with_duplicate_sources():
    # Multiple edges with same source
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}]
    # Both 2 and 3 are sinks, should return 2 (first in nodes)
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.92μs -> 1.25μs (53.4% faster)


def test_nodes_with_non_integer_ids():
    # 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.21μs (62.2% faster)


def test_nodes_with_mixed_id_types():
    # IDs are mixed types (should work as long as equality holds)
    nodes = [{"id": 1}, {"id": "2"}]
    edges = [{"source": 1, "target": "2"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.00μs -> 1.17μs (71.4% faster)


def test_multiple_sinks_with_complex_node():
    # Node dicts with extra data
    nodes = [
        {"id": 1, "name": "A"},
        {"id": 2, "name": "B"},
        {"id": 3, "name": "C"},
    ]
    edges = [{"source": 1, "target": 2}]
    # Both 2 and 3 are sinks, should return 2 (first in nodes)
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.17μs (57.2% faster)


def test_edges_with_extra_fields():
    # Edges have 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.21μs (51.7% faster)


def test_nodes_with_duplicate_ids():
    # Duplicate node IDs: function returns the first sink node found
    nodes = [{"id": 1}, {"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    # Both node 1s have outgoing edges, node 2 is sink
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.00μs -> 1.25μs (60.0% faster)


def test_edge_with_missing_source_key():
    # Edge missing 'source' key (should raise KeyError)
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"target": 2}]
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 1.62μs -> 792ns (105% faster)


def test_large_linear_chain():
    # Large chain: 1->2->3->...->1000, should return node 1000
    N = 1000
    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.2ms -> 62.8μs (28875% faster)


def test_large_star_graph():
    # Star graph: 1->2, 1->3, ..., 1->1000; sinks are 2..1000, should return 2
    N = 1000
    nodes = [{"id": i} for i in range(1, N + 1)]
    edges = [{"source": 1, "target": i} for i in range(2, N + 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 37.7μs -> 20.3μs (85.2% faster)


def test_large_no_edges():
    # 1000 nodes, no edges: should return first node
    N = 1000
    nodes = [{"id": i} for i in range(1, N + 1)]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.42μs -> 1.08μs (30.7% faster)


def test_large_all_nodes_have_outgoing_edges():
    # 1000 nodes, each has outgoing edge to next (cycle): should return None
    N = 1000
    nodes = [{"id": i} for i in range(1, N + 1)]
    edges = [{"source": i, "target": (i % N) + 1} for i in range(1, N + 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.3ms -> 62.2μs (29318% faster)


def test_large_multiple_sinks():
    # 1000 nodes, edges from 1->2, 1->3, ..., 1->1000; sinks are 2..1000, should return 2
    N = 1000
    nodes = [{"id": i} for i in range(1, N + 1)]
    edges = [{"source": 1, "target": i} for i in range(2, N + 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 37.8μs -> 20.4μs (85.1% 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-mjnd1wp9 and push.