⚡️ Speed up function find_last_node by 24,006%

[codeflash-ai]

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

⏱️ Runtime : 63.5 milliseconds → 263 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 240x speedup by eliminating a quadratic time complexity bottleneck through a simple but powerful algorithmic change.

Key Optimization:
The original implementation uses a nested loop structure that checks all(e["source"] != n["id"] for e in edges) for each node. This creates O(n × m) complexity where n is the number of nodes and m is the number of edges. For each node, it iterates through all edges to verify the node isn't a source.

The optimized version pre-computes a set of all source node IDs with sources = {e["source"] for e in edges}, then performs O(1) set membership lookups with n["id"] not in sources. This reduces the overall complexity to O(n + m).

Why This Matters:

• Set lookups are O(1) vs iterating through all edges which is O(m)
• The optimization is most impactful when there are many edges (line profiler shows 493ms → 0.414ms total time)
• Large-scale tests demonstrate the dramatic improvement:
  • test_large_linear_chain (1000 nodes): 18.1ms → 55.6μs (324x faster)
  • test_large_almost_circular_graph (999 nodes): 18.1ms → 55.1μs (327x faster)
  • test_large_graph_with_multiple_leaves (1000 nodes): 4.50ms → 28.2μs (158x faster)

Performance Characteristics:

• Small graphs (2-5 nodes): 50-94% faster - modest but consistent gains
• Medium graphs (63 nodes): 10x faster
• Large graphs (500-1000 nodes): 150-325x faster - scales linearly instead of quadratically

Edge Cases:
The optimization maintains correctness across all scenarios including empty inputs, cycles, disconnected graphs, self-loops, and non-integer IDs. There's a slight regression (5-22% slower) on truly empty graphs where the set creation overhead isn't amortized, but this is negligible in absolute terms (nanoseconds).

This optimization is particularly valuable if find_last_node is called repeatedly on graphs with many edges, as the algorithmic improvement compounds with scale.

✅ 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():
    # One node, no edges; should return the node itself
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.17μs -> 958ns (21.8% faster)


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


def test_three_nodes_linear():
    # Three nodes in a line: 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.21μs -> 1.25μs (76.6% faster)


def test_three_nodes_branching():
    # Three nodes, 1 branches to 2 and 3; both 2 and 3 are leaves, should return first 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.75μs -> 1.17μs (50.0% faster)


def test_multiple_nodes_multiple_leaves():
    # Four nodes, 1->2, 1->3, 2->4; leaves are 3 and 4
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}]
    edges = [
        {"source": 1, "target": 2},
        {"source": 1, "target": 3},
        {"source": 2, "target": 4},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.25μs -> 1.25μs (80.0% 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  # 750ns -> 834ns (10.1% slower)


def test_nodes_with_no_edges():
    # Multiple nodes, but no edges; all are leaves, should return first node
    nodes = [{"id": "a"}, {"id": "b"}, {"id": "c"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.12μs -> 1.00μs (12.5% faster)


def test_circular_graph():
    # Circular graph: 1->2, 2->3, 3->1; no leaves, 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.12μs -> 1.29μs (64.6% faster)


def test_disconnected_graph():
    # Disconnected nodes: 1->2, node 3 is isolated
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.75μs -> 1.17μs (50.1% faster)


def test_node_with_self_loop():
    # Node 1 has a self-loop; should not be a leaf
    nodes = [{"id": 1}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.33μs -> 1.08μs (23.1% faster)


def test_multiple_edges_from_one_node():
    # Node 1 points to 2 and 3, 2 points to 4, 3 points to 5; leaves are 4 and 5
    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.67μs -> 1.38μs (94.0% faster)


def test_non_integer_node_ids():
    # Node ids are strings
    nodes = [{"id": "x"}, {"id": "y"}]
    edges = [{"source": "x", "target": "y"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.17μs (60.7% 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.75μs -> 1.12μs (55.6% faster)


def test_nodes_with_extra_fields():
    # Nodes have extra fields, should be returned as is
    nodes = [{"id": 1, "label": "start"}, {"id": 2, "label": "end"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.75μs -> 1.12μs (55.6% faster)


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


def test_large_linear_chain():
    # 1000 nodes in a line: 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)
    result = codeflash_output  # 18.1ms -> 55.6μs (32445% faster)


def test_large_branching_tree():
    # Binary tree up to depth 5 (max 63 nodes), leaves are last 32 nodes
    nodes = [{"id": i} for i in range(63)]
    edges = []
    for i in range(31):  # Each node has two children except leaves
        edges.append({"source": i, "target": 2 * i + 1})
        edges.append({"source": i, "target": 2 * i + 2})
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 45.1μs -> 3.79μs (1090% faster)


def test_large_disconnected_graph():
    # 1000 nodes, no edges; all are leaves, should return first node
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.29μs -> 1.00μs (29.1% faster)


def test_large_almost_circular_graph():
    # 999 nodes in a cycle, 1 node (id=999) is isolated; should return node 999
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [
        {"source": i, "target": (i + 1) % (N - 1)} for i in range(N - 1)
    ]  # 0->1->2->...->998->0
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.1ms -> 55.1μs (32703% faster)


def test_large_graph_with_multiple_leaves():
    # 500 nodes, each node i points to i+1, plus 500 isolated nodes
    N = 500
    nodes = [{"id": i} for i in range(N)] + [{"id": f"x{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  # 4.50ms -> 28.2μs (15843% faster)
    # Leaves are node 499 and all x0..x499; function returns first found
    possible_leaves = [{"id": N - 1}] + [{"id": f"x{i}"} for i in range(N)]


# 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; node should be returned as the last node
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.17μs -> 958ns (21.8% faster)


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


def test_three_nodes_linear_chain():
    # 1 -> 2 -> 3; node 3 should be 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.17μs -> 1.25μs (73.4% faster)


def test_multiple_possible_last_nodes():
    # 1 -> 2, 1 -> 3; both 2 and 3 have no outgoing edges, should return 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.83μs -> 1.17μs (57.1% faster)


def test_no_edges_multiple_nodes():
    # All nodes have no outgoing edges; should return the first node
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.21μs -> 958ns (26.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_but_no_nodes():
    # Edges exist, but no nodes; should return None
    nodes = []
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 750ns -> 958ns (21.7% slower)


def test_cycle_graph():
    # 1 -> 2 -> 3 -> 1; 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.12μs -> 1.29μs (64.5% faster)


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


def test_disconnected_nodes():
    # Some nodes not connected at all; should return first disconnected node as last node
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.75μs -> 1.12μs (55.6% faster)
    # Let's check for 3 as well
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    result_ids = [n["id"] for n in nodes if all(e["source"] != n["id"] for e in edges)]


def test_non_integer_node_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.42μs -> 1.21μs (99.9% faster)


def test_duplicate_node_ids():
    # Duplicate node IDs in nodes list; should return the first one not a source
    nodes = [{"id": 1}, {"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_missing_source_in_edges():
    # Edge references a source not in nodes; should not affect last node selection
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 3, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.42μs -> 1.04μs (36.0% faster)


def test_missing_target_in_edges():
    # Edge references a target not in nodes; should not affect last node selection
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.71μs -> 1.12μs (51.9% faster)


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


def test_large_linear_chain():
    # Large linear chain: 1 -> 2 -> ... -> 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 -> 55.8μs (32511% faster)


def test_large_branching_tree():
    # Binary tree, depth = 5 (31 nodes), all leaves are last nodes
    from collections import deque

    nodes = [{"id": i} for i in range(1, 32)]
    edges = []
    for i in range(1, 16):  # non-leaf nodes
        edges.append({"source": i, "target": 2 * i})
        edges.append({"source": i, "target": 2 * i + 1})
    # The first leaf node is 16
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 14.5μs -> 2.29μs (531% faster)


def test_large_graph_all_sources():
    # All nodes are sources (cycle), should return None
    n = 500
    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  # 4.51ms -> 28.0μs (15987% faster)


def test_large_graph_all_leaves():
    # All nodes have no outgoing edges; should return first node
    n = 1000
    nodes = [{"id": i} for i in range(n)]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.29μs -> 1.04μs (24.1% faster)


def test_large_graph_sparse_edges():
    # Only a few edges in a large node set
    n = 1000
    nodes = [{"id": i} for i in range(n)]
    edges = [{"source": 0, "target": 1}, {"source": 2, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.92μs -> 1.29μs (48.5% 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-mjj66s8z and push.