⚡️ Speed up function find_last_node by 34,233%

[codeflash-ai]

📄 34,233% (342.33x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 178 milliseconds → 520 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 342x speedup by replacing a nested-loop O(N×M) algorithm with a linear O(N+M) approach using a set for membership testing.

Key Optimization:

The original code uses all(e["source"] != n["id"] for e in edges) for each node, checking every edge repeatedly. For N nodes and M edges, this results in N×M comparisons—catastrophic for large graphs.

The optimized version:

1. Pre-builds a set of all source IDs once: edge_source_ids = {e["source"] for e in edges} (O(M))
2. Uses O(1) set membership instead of O(M) linear scan: n["id"] not in edge_source_ids
3. Adds early-exit guard for empty nodes list

Why This Works:

Python sets provide O(1) average-case lookup via hash tables, versus O(M) for iterating through all edges. For each node, the cost drops from M comparisons to a single hash lookup.

Impact by Test Case:

• Small graphs (2-3 nodes): 132-221% faster—modest gain since overhead is small
• Linear chains (1000 nodes): 57,566-65,412% faster—the nested loop penalty becomes extreme as both N and M grow
• Dense graphs (100 nodes, 9,900 edges): 17,131% faster—the original's O(N×M) ≈ 990,000 operations vs. optimized O(N+M) ≈ 10,000
• Star graphs (1 source, 999 targets): 210% faster—M is large but only one source, so set-building still helps

Workload Suitability:

Without function_references, we can't confirm hot-path usage, but the optimization is universally beneficial: it maintains correctness while drastically reducing algorithmic complexity. Any workload with moderate-to-large graphs (>100 nodes/edges) will see substantial gains.

✅ 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

# ----------------
# 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)  # 2.71μs -> 1.00μs (171% faster)


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


def test_three_nodes_linear_chain():
    # 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)  # 3.88μs -> 1.21μs (221% faster)


def test_multiple_possible_last_nodes():
    # 1->2, 3 (no edges): both 2 and 3 have no outgoing edges, but 2 appears first
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    # Should return the first node in nodes with no outgoing edges (2)
    codeflash_output = find_last_node(nodes, edges)  # 2.83μs -> 1.17μs (143% faster)


def test_multiple_last_nodes_order():
    # 1->2, 3 (no edges): but 3 appears before 2 in nodes
    nodes = [{"id": 1}, {"id": 3}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    # Should return 3, as it is first in nodes with no outgoing edges
    codeflash_output = find_last_node(nodes, edges)  # 2.71μs -> 1.17μs (132% faster)


# ----------------
# Edge Test Cases
# ----------------


def test_empty_nodes_and_edges():
    # No nodes, no edges: should return None
    codeflash_output = find_last_node([], [])  # 1.21μs -> 250ns (383% faster)


def test_nodes_no_edges():
    # Multiple nodes, no edges: should return the first node
    nodes = [{"id": "a"}, {"id": "b"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 2.12μs -> 1.04μs (104% faster)


def test_all_nodes_have_outgoing_edges():
    # 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)  # 3.67μs -> 1.33μs (175% 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}]
    # 2 has no outgoing edges, so should be returned
    codeflash_output = find_last_node(nodes, edges)  # 2.75μs -> 1.17μs (136% faster)


def test_edges_with_nonexistent_nodes():
    # Edges refer to non-existent nodes: should not affect result
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 3, "target": 1}, {"source": 1, "target": 2}]
    # 2 has no outgoing edges
    codeflash_output = find_last_node(nodes, edges)  # 3.04μs -> 1.17μs (161% faster)


def test_duplicate_edges():
    # Duplicate 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},
    ]
    # 3 has no outgoing edges
    codeflash_output = find_last_node(nodes, edges)  # 3.62μs -> 1.33μs (172% faster)


def test_node_with_non_integer_id():
    # Node IDs can be strings or other hashables
    nodes = [{"id": "a"}, {"id": "b"}]
    edges = [{"source": "a", "target": "b"}]
    codeflash_output = find_last_node(nodes, edges)  # 3.12μs -> 1.17μs (168% faster)


def test_nodes_with_extra_fields():
    # Nodes may have extra fields, should still return the full node
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)  # 2.71μs -> 1.12μs (141% faster)


def test_edge_with_extra_fields():
    # Edges may have extra fields, should not affect result
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2, "weight": 3}]
    codeflash_output = find_last_node(nodes, edges)  # 2.83μs -> 1.17μs (143% faster)


def test_multiple_edges_from_one_node():
    # One node with multiple outgoing edges
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}]
    # 2 is first in nodes with no outgoing edges
    codeflash_output = find_last_node(nodes, edges)  # 2.92μs -> 1.21μs (141% faster)


def test_node_with_incoming_but_no_outgoing():
    # Node with only incoming edges is a last node
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 3}, {"source": 2, "target": 3}]
    # 1 and 2 have outgoing edges, 3 does not
    codeflash_output = find_last_node(nodes, edges)  # 3.38μs -> 1.21μs (179% faster)


def test_nodes_with_none_id():
    # Node id is None
    nodes = [{"id": None}, {"id": 2}]
    edges = [{"source": 2, "target": None}]
    # None has no outgoing edges, so should be returned
    codeflash_output = find_last_node(nodes, edges)  # 2.50μs -> 1.08μs (131% faster)


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


def test_large_linear_chain():
    # Large linear chain of 1000 nodes
    nodes = [{"id": i} for i in range(1000)]
    edges = [{"source": i, "target": i + 1} for i in range(999)]
    # Last node is 999
    codeflash_output = find_last_node(nodes, edges)  # 31.1ms -> 53.9μs (57566% faster)


def test_large_star_graph():
    # Star graph: node 0 points to all others, others have no outgoing edges
    nodes = [{"id": i} for i in range(1000)]
    edges = [{"source": 0, "target": i} for i in range(1, 1000)]
    # First node with no outgoing edges is 1
    codeflash_output = find_last_node(nodes, edges)  # 61.3μs -> 19.8μs (211% faster)


def test_large_all_edges():
    # Every node points to every other node (no last node)
    nodes = [{"id": i} for i in range(100)]
    edges = [
        {"source": i, "target": j} for i in range(100) for j in range(100) if i != j
    ]
    codeflash_output = find_last_node(nodes, edges)  # 32.9ms -> 190μs (17131% faster)


def test_large_sparse_graph_multiple_last_nodes():
    # Many nodes, only a few edges, multiple last nodes
    nodes = [{"id": i} for i in range(1000)]
    edges = [{"source": 0, "target": 1}, {"source": 2, "target": 3}]
    # First node with no outgoing edges is 1
    codeflash_output = find_last_node(nodes, edges)  # 4.29μs -> 1.29μs (232% faster)


def test_large_graph_with_non_sequential_ids():
    # Nodes with non-sequential IDs
    nodes = [{"id": i * 10} for i in range(1000)]
    edges = [{"source": i * 10, "target": (i + 1) * 10} for i in range(999)]
    codeflash_output = find_last_node(nodes, edges)  # 37.2ms -> 61.4μs (60536% 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
    nodes = [{"id": 1, "name": "A"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.96μs -> 1.04μs (184% faster)


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


def test_three_nodes_linear_chain():
    # Three nodes, edges 1->2, 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  # 4.00μs -> 1.25μs (220% faster)


def test_three_nodes_multiple_leaves():
    # Three nodes, edges 1->2 and 1->3: should return node 2 (first found leaf)
    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  # 3.04μs -> 1.21μs (152% faster)


def test_node_with_multiple_incoming_edges():
    # Node 3 has two incoming edges, but no outgoing: should return node 3
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 3}, {"source": 2, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 3.50μs -> 1.25μs (180% faster)


# -----------------------
# Edge Test Cases
# -----------------------


def test_no_nodes():
    # No nodes: should return None
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.33μs -> 250ns (433% faster)


def test_no_edges_multiple_nodes():
    # 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  # 2.29μs -> 959ns (139% faster)


def test_all_nodes_have_outgoing_edges():
    # All nodes have outgoing edges: 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  # 2.92μs -> 1.21μs (141% faster)


def test_cycle_graph():
    # Cycle: 1->2->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  # 3.71μs -> 1.29μs (187% 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)
    result = codeflash_output  # 3.38μs -> 1.17μs (189% faster)


def test_node_with_extra_fields():
    # Node dicts with extra fields
    nodes = [{"id": 1, "data": "foo"}, {"id": 2, "data": "bar"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.88μs -> 1.21μs (138% faster)


def test_disconnected_graph():
    # Disconnected graph: two components, should return first leaf in nodes order
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}]
    edges = [{"source": 1, "target": 2}, {"source": 3, "target": 4}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 3.04μs -> 1.21μs (152% faster)


def test_edge_with_nonexistent_node():
    # Edge references a node not in nodes; should still return the correct leaf
    nodes = [{"id": 1}, {"id": 2}]
    edges = [
        {"source": 1, "target": 2},
        {"source": 3, "target": 1},
    ]  # source 3 not in nodes
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 3.04μs -> 1.21μs (152% faster)


def test_nodes_out_of_order():
    # Nodes not sorted by id, function should return first leaf in nodes order
    nodes = [{"id": 10}, {"id": 3}, {"id": 7}]
    edges = [{"source": 10, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.79μs -> 1.17μs (139% faster)


def test_duplicate_node_ids():
    # Duplicate node ids: should return the first leaf (by nodes order) not referenced as a source
    nodes = [{"id": 1}, {"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 3.17μs -> 1.21μs (162% faster)


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


def test_large_linear_chain():
    # Large chain of 1000 nodes: should return the last node
    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  # 36.0ms -> 55.0μs (65412% faster)


def test_large_star_graph():
    # Star graph: node 0 points to all others, leaves are 1..999, should return node 1 (first leaf)
    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  # 61.4μs -> 19.8μs (210% faster)


def test_large_no_edges():
    # 1000 nodes, no 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  # 2.62μs -> 1.00μs (162% faster)


def test_large_all_nodes_have_outgoing():
    # All nodes have outgoing edges, forming a cycle: should return None
    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  # 33.0ms -> 53.7μs (61417% faster)


def test_large_multiple_leaves():
    # 1000 nodes, first 500 are sources, last 500 are leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i + 500} for i in range(500)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 7.94ms -> 27.3μs (28940% 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-mjiyqvy6 and push.