⚡️ Speed up function find_last_node by 22,046%

[codeflash-ai]

📄 22,046% (220.46x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 78.4 milliseconds → 354 microseconds (best of 250 runs)

📝 Explanation and details

The optimization achieves a 220x speedup by replacing a nested loop with a set-based lookup, dramatically reducing algorithmic complexity from O(n*m) to O(n+m) where n is the number of nodes and m is the number of edges.

Key Changes:

1. Pre-compute edge sources as a set: Instead of iterating through all edges for each node (nested loop), the optimized code builds a set of all source node IDs upfront: edge_sources = {e["source"] for e in edges}. This single pass through edges takes O(m) time.

2. O(1) membership test: Set membership (n["id"] not in edge_sources) is O(1) average case, versus the original's O(m) check using all(e["source"] != n["id"] for e in edges) which had to compare against every edge.

3. Edge case handling: The code explicitly handles the empty edges case to match original behavior - when there are no edges, it returns the first node without attempting to access n["id"], avoiding potential KeyErrors on malformed input.

Why It's Faster:

The original code had quadratic-like behavior: for each node, it scanned all edges. With 1000 nodes and 999 edges (linear chain), this meant ~1 million comparisons. The optimized version does 999 edge scans + 1000 set lookups = ~2000 operations total.

Performance Characteristics:

• Small graphs (2-10 nodes): 40-60% faster - set construction overhead is minimal
• Medium graphs (100+ nodes): 80-120% faster - set lookup benefits become clear
• Large graphs (1000 nodes): 330x faster on linear chains, 336x on cycles - the O(n*m) penalty dominates in original code
• Disconnected nodes (no edges): 40-50% faster - early return path avoids unnecessary checks
• Empty inputs: 5-30% slower - added conditional check overhead, but negligible in absolute terms (nanoseconds)

This optimization is particularly valuable when find_last_node is called frequently or on larger graphs, as the speedup scales superlinearly with input size.

✅ Correctness verification report:

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

🌀 Click to see Generated Regression Tests

import pytest
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.21μs -> 833ns (45.1% 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)
    result = codeflash_output  # 1.83μs -> 1.21μs (51.7% faster)


def test_three_nodes_linear_chain():
    # Three nodes, edges 1->2, 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.08μs -> 1.29μs (61.3% faster)


def test_multiple_last_nodes():
    # Two separate chains: 1->2 and 3->4, so last nodes are 2 and 4, should return 2 (first found)
    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  # 1.79μs -> 1.25μs (43.4% faster)


def test_no_edges_multiple_nodes():
    # Multiple nodes, no edges: should return the first node
    nodes = [{"id": 10}, {"id": 20}, {"id": 30}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.17μs -> 833ns (40.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  # 750ns -> 792ns (5.30% slower)


def test_edges_but_no_nodes():
    # Edges exist but nodes list is empty: should return None
    nodes = []
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 750ns -> 1.08μs (30.7% slower)


def test_edges_with_nonexistent_nodes():
    # Edges refer to node ids not in nodes list: should return first node
    nodes = [{"id": 100}]
    edges = [{"source": 999, "target": 1000}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.46μs -> 1.12μs (29.6% faster)


def test_cycle_graph():
    # Cycle: 1->2, 2->3, 3->1, 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.17μs -> 1.33μs (62.6% faster)


def test_node_with_multiple_incoming_edges():
    # 1->3, 2->3: last node is 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  # 2.08μs -> 1.29μs (61.3% faster)


def test_node_with_multiple_outgoing_edges():
    # 1->2, 1->3, 2->4, 3->4: last node is 4
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}]
    edges = [
        {"source": 1, "target": 2},
        {"source": 1, "target": 3},
        {"source": 2, "target": 4},
        {"source": 3, "target": 4},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.58μs -> 1.42μs (82.3% faster)


def test_node_with_self_loop():
    # 1->1 (self loop): should return None
    nodes = [{"id": 1}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.33μs -> 1.12μs (18.6% faster)


def test_duplicate_edges():
    # Duplicate edges 1->2, 1->2: last node is 2
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.79μs -> 1.29μs (38.7% faster)


def test_nodes_with_extra_fields():
    # Nodes with extra fields, should still work
    nodes = [{"id": 1, "data": "foo"}, {"id": 2, "data": "bar"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.75μs -> 1.21μs (44.9% faster)


def test_edges_with_extra_fields():
    # Edges with extra fields, should ignore them
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2, "weight": 5}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.75μs -> 1.17μs (50.0% faster)


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


def test_large_linear_chain():
    # Large chain of 1000 nodes: last node is 1000
    nodes = [{"id": i} for i in range(1, 1001)]
    edges = [{"source": i, "target": i + 1} for i in range(1, 1000)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.3ms -> 54.9μs (33269% faster)


def test_large_star_graph():
    # Star graph: 1->2, 1->3, ..., 1->1000. All nodes except 1 are last nodes, should return 2 (first found)
    nodes = [{"id": i} for i in range(1, 1001)]
    edges = [{"source": 1, "target": i} for i in range(2, 1001)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 37.8μs -> 20.8μs (82.1% faster)


def test_large_disconnected_nodes():
    # 1000 nodes, no edges: should return first node
    nodes = [{"id": i} for i in range(1, 1001)]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.33μs -> 875ns (52.5% faster)


def test_large_cycle():
    # 1000 node cycle: 1->2, ..., 1000->1, should return None
    nodes = [{"id": i} for i in range(1, 1001)]
    edges = [{"source": i, "target": i % 1000 + 1} for i in range(1, 1001)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.5ms -> 54.9μs (33524% faster)


def test_large_multiple_chains():
    # Two chains: 1->2->...->500 and 501->502->...->1000, so last nodes are 500 and 1000
    nodes = [{"id": i} for i in range(1, 1001)]
    edges = [{"source": i, "target": i + 1} for i in range(1, 500)] + [
        {"source": i, "target": i + 1} for i in range(501, 1000)
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 4.56ms -> 37.5μs (12055% 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

# 1. BASIC TEST CASES


def test_single_node_no_edges():
    # Single node, no edges: node is the last node
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.21μs -> 833ns (45.1% faster)


def test_two_nodes_one_edge():
    # Two nodes, one edge from 1 to 2: node 2 is the last node
    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_three_nodes_linear_chain():
    # 1->2->3: node 3 is the 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.25μs -> 1.29μs (74.3% faster)


def test_multiple_end_nodes():
    # 1->2, 1->3: nodes 2 and 3 are both last nodes, returns the first in nodes list
    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.25μs (46.6% faster)


def test_no_edges_multiple_nodes():
    # All nodes are last nodes, returns the first in nodes list
    nodes = [{"id": 10}, {"id": 20}, {"id": 30}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.21μs -> 833ns (45.1% faster)


# 2. EDGE TEST CASES


def test_empty_nodes():
    # No nodes: should return None
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 750ns -> 833ns (9.96% slower)


def test_edges_with_no_matching_nodes():
    # Edges reference node ids not present in nodes
    nodes = [{"id": 1}]
    edges = [{"source": 2, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.42μs -> 1.12μs (26.0% faster)


def test_node_with_self_loop():
    # Node points to itself, should not be considered a last node
    nodes = [{"id": 1}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.33μs -> 1.12μs (18.5% faster)


def test_cycle_graph():
    # 1->2->3->1 (cycle): all nodes are sources, so there is 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)
    result = codeflash_output  # 2.21μs -> 1.29μs (71.0% faster)


def test_duplicate_edges():
    # 1->2 (twice): node 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)
    result = codeflash_output  # 1.83μs -> 1.29μs (41.9% faster)


def test_non_integer_node_ids():
    # Node 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.25μs (56.6% faster)


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


def test_edge_with_extra_keys():
    # Edge has extra keys, should not affect result
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2, "weight": 5}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.17μs (60.7% faster)


def test_nodes_with_duplicate_ids():
    # Nodes with duplicate ids: function should return the first matching last node
    nodes = [{"id": 1}, {"id": 2}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.75μs -> 1.17μs (50.0% faster)


# 3. 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.4ms -> 54.8μs (33538% faster)


def test_large_star_graph():
    # Star graph: 1->2, 1->3, ..., 1->1000
    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  # 38.3μs -> 21.0μs (82.2% faster)


def test_large_no_edges():
    # 1000 nodes, no edges: returns 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.33μs -> 959ns (39.0% faster)


def test_large_all_sources():
    # All nodes are sources in at least one edge: no last node
    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)
    ]  # forms a cycle
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.5ms -> 54.8μs (33680% faster)


def test_large_multiple_last_nodes():
    # 1->2, 1->3, 1->4, ..., 1->1000: nodes 2..1000 are all last nodes, should return node 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.9μs -> 21.0μs (80.9% 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-mjj73yso and push.