⚡️ Speed up function find_last_node by 16,863%

[codeflash-ai]

📄 16,863% (168.63x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 52.0 milliseconds → 307 microseconds (best of 152 runs)

📝 Explanation and details

Brief: The optimized version replaces a repeated scan of the edges for every node with a single pass that collects all edge sources into a set, turning an O(N*M) check into O(N+M) work. It also adds a small fast-path for the empty-edges case to avoid unnecessary work and preserve the original function's behavior. These changes explain the ~168x measured speedup (52 ms → 307 μs).

What changed and why it’s faster

• Precompute sources set: The original code used
  next((n for n in nodes if all(e["source"] != n["id"] for e in edges)), None)
  which, for each node, iterates over edges (the all(...) generator). That is O(N * M) comparisons in the worst case (N = number of nodes, M = number of edges).
  The optimized code computes sources = {e["source"] for e in edges} once (O(M)) and then checks n["id"] not in sources for each node (O(1) average per membership test). Total complexity becomes O(M + N).
• Fast-path for empty edges: If edges is empty, the original all(...) check is vacuously true and returns the first node. The optimized code preserves that behavior with if not edges: return next(iter(nodes), None). This avoids building an unnecessary set and is faster for the common "no edges" case.
• Small memory/time tradeoff: We allocate a set of unique sources (size ≤ M). The small extra memory is offset by the large reduction in repeated iteration when M and N are non-trivial.

Evidence in the profiler

• Original line profiler shows all time was spent in the single generator/all(...) check (repeated traversal of edges per node).
• Optimized profiler shows time split between building the sources set and doing cheap membership tests. Building the set is linear in edges and then each node check is a single O(1) membership test — much cheaper than re-scanning edges each time.

Behavioral impact and correctness

• Semantics preserved:
  • Returns the first node that is not a source in any edge (same as original).
  • When edges is empty, returns the first node (preserved by the fast-path).
  • Malformed edges missing "source" still raise KeyError (behavior unchanged).
  • Duplicate ids and type-mismatch behavior remains identical.
• Performance tradeoffs:
  • For very small inputs (e.g., empty nodes or tiny lists), the micro-overhead of branching and set creation can be similar or slightly higher; tests show one tiny regression when nodes are empty (~15% slower in one case). This is expected and negligible in real workloads.
  • For moderate-to-large inputs (the heavy cases), the improvement is dramatic — tests like large_chain and large_unordered_edges show orders-of-magnitude speedups.

When this matters (based on tests)

• Best for workloads with many nodes and/or many edges (large_chain, large_unordered_edges): huge wins because you avoid repeated scanning of edges.
• Also benefits normal cases like duplicate ids, unordered edges, and typical graphs — basically any non-trivial graph.
• Minimal downside for trivial inputs.

Complexity summary

• Original: O(N * M) time, O(1) extra space.
• Optimized: O(N + M) time, O(M) extra space (for the set of sources).

In short: change from repeated edge scans to a single set-construction + O(1) membership checks explains the significant runtime improvement while preserving the original function’s behavior.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 22 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 tests


def test_basic_chain_returns_terminal_node():
    # Simple chain: 1 -> 2 -> 3. Node with id "3" has no outgoing edges, so it 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  # 3.58μs -> 1.71μs (110% faster)


def test_empty_edges_returns_first_node():
    # When there are no edges, every node has no outgoing edge.
    # The implementation should return the first node in the nodes iterable.
    nodes = [{"id": "alpha"}, {"id": "beta"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 667ns (181% faster)


def test_empty_nodes_returns_none():
    # If there are no nodes, function should return None.
    nodes = []
    edges = [{"source": "x", "target": "y"}]
    codeflash_output = find_last_node(nodes, edges)  # 1.17μs -> 1.38μs (15.2% slower)


# Edge tests


def test_non_string_ids_work_correctly():
    # IDs are integers instead of strings. The comparison should still work.
    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  # 3.54μs -> 2.00μs (77.1% faster)


def test_duplicate_node_ids_selects_first_matching_node():
    # If multiple nodes share the same id and that id is not used as a source,
    # the function should return the first such node encountered.
    nodes = [{"id": "dup", "value": 1}, {"id": "dup", "value": 2}, {"id": "unique"}]
    edges = [{"source": "unique", "target": "x"}]  # 'dup' is not a source in edges
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 1.46μs (34.4% faster)


def test_multiple_nodes_without_outgoing_edges_returns_first_of_them():
    # When more than one node has no outgoing edges, the first such node in 'nodes'
    # should be returned (due to 'next' over nodes).
    nodes = [{"id": "a"}, {"id": "b"}, {"id": "c"}]
    edges = [{"source": "a", "target": "b"}]  # 'b' and 'c' have no outgoing edges
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.50μs -> 1.54μs (62.2% faster)


def test_malformed_edge_missing_source_raises_key_error():
    # If an edge dict doesn't have 'source', accessing e['source'] raises KeyError.
    nodes = [{"id": "only"}]
    edges = [{"target": "only"}]  # missing 'source'
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)  # 4.08μs -> 1.46μs (180% faster)


def test_returned_node_is_same_object_as_in_input_list():
    # Ensure the function returns the very object from the nodes list (not a copy).
    nodes = [{"id": "n1", "flag": False}, {"id": "n2"}]
    edges = [
        {"source": "n2", "target": "x"}
    ]  # only n2 has an outgoing edge -> n1 has none
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.38μs -> 1.58μs (49.9% faster)
    # Modify the returned object and check the original list reflects change
    result["flag"] = True


def test_large_chain_of_nodes_returns_last_element():
    # Build a chain of 500 nodes (0 -> 1 -> 2 -> ... -> 499).
    # The last node (id "499") has no outgoing edge and should be returned.
    size = 500  # well under 1000
    nodes = [{"id": str(i)} for i in range(size)]
    # Create edges chaining 0->1, 1->2, ..., 498->499
    edges = [{"source": str(i), "target": str(i + 1)} for i in range(size - 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 9.84ms -> 71.9μs (13592% faster)


def test_large_unordered_edges_and_nodes_still_identifies_last_node():
    # Create 800 nodes and many edges in shuffled-like order to ensure
    # re-checking edges for each node works correctly.
    size = 800  # keep data structures under 1000
    nodes = [{"id": f"n{i}"} for i in range(size)]
    edges = []
    # For a more complex edge set, add many edges including self-loops and additional cross links,
    # but ensure only the last node 'n{size-1}' has no outgoing edge.
    for i in range(size - 1):
        # primary chain edge
        edges.append({"source": f"n{i}", "target": f"n{i+1}"})
        # add a couple of extra edges to earlier nodes to increase noise
        if i % 10 == 0:
            edges.append({"source": f"n{i}", "target": f"n0"})
        if i % 15 == 0:
            edges.append({"source": f"n{i}", "target": f"n{(i+3) % (size-1)}"})
    # There should be many edges but only node 'n{size-1}' has no outgoing edge.
    codeflash_output = find_last_node(nodes, edges)
    res = codeflash_output  # 32.6ms -> 132μs (24477% faster)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from typing import Any, Dict, List

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

# unit tests


def test_basic_chain_returns_last_node_identity():
    # Basic scenario: a simple chain 1->2->3, and a standalone node 4
    # Expectation: node with id "4" has no outgoing edges and should be returned.
    nodes = [{"id": "1"}, {"id": "2"}, {"id": "3"}, {"id": "4"}]
    edges = [{"source": "1", "target": "2"}, {"source": "2", "target": "3"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 3.79μs -> 1.88μs (102% faster)


def test_multiple_sinks_returns_first_in_nodes_order():
    # If multiple nodes have no outgoing edges, the function should return
    # the first of those nodes as they appear in the nodes list.
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    # Only A has outgoing edge; B and C are sinks; B appears first among sinks
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.75μs -> 1.62μs (69.2% faster)


def test_empty_nodes_returns_none():
    # Edge case: no nodes at all -> should return None
    nodes: List[Dict[str, Any]] = []
    edges = [{"source": "1", "target": "2"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 1.42μs (11.8% slower)


def test_empty_edges_returns_first_node():
    # Edge case: empty edges means no node has outgoing edges -> first node returned
    nodes = [{"id": "x"}, {"id": "y"}]
    edges: List[Dict[str, Any]] = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.00μs -> 583ns (243% faster)


def test_edges_reference_external_nodes_ignored():
    # Edges may reference node ids not present in nodes list; those edges
    # should not prevent local nodes from being considered sinks.
    nodes = [{"id": "N1"}, {"id": "N2"}]
    # Edge references external node "EXT" and has a source "EXT" that doesn't match any node
    edges = [{"source": "EXT", "target": "ZZ"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.04μs -> 1.54μs (32.4% faster)


def test_malformed_edge_raises_key_error():
    # If an edge dict is malformed (missing "source"), the function accesses e["source"]
    # and should raise a KeyError. This test documents that behavior.
    nodes = [{"id": "1"}]
    edges = [{"target": "1"}]  # missing "source" -> KeyError expected
    with pytest.raises(KeyError):
        codeflash_output = find_last_node(nodes, edges)
        _ = codeflash_output  # 4.08μs -> 1.25μs (227% faster)


def test_duplicate_node_ids_behavior():
    # When nodes contain duplicate ids, the function returns the first node in list
    # that has no outgoing edges. If edges contain a source equal to that id, then
    # none of the nodes with that id are sinks.
    # Case A: duplicates with no outgoing edge for that id -> first duplicate returned
    nodes = [{"id": "dup"}, {"id": "dup"}, {"id": "other"}]
    edges = [{"source": "other", "target": "dup"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.75μs -> 1.58μs (73.7% faster)

    # Case B: duplicates but edges include source "dup" -> both duplicates are not sinks
    edges_with_dup_source = [{"source": "dup", "target": "other"}]
    codeflash_output = find_last_node(nodes, edges_with_dup_source)
    result2 = codeflash_output  # 2.08μs -> 958ns (117% faster)


def test_cycle_returns_none():
    # A cycle where every node has an outgoing edge should yield None
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "A"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.79μs -> 1.75μs (59.5% faster)


def test_type_mismatch_between_ids_and_sources():
    # If node ids and edge sources have different types (e.g., one str, one int),
    # they should not match. This verifies strict equality behavior.
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": "1", "target": "2"}]  # sources are strings, nodes ids are ints
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.25μs -> 1.54μs (45.9% faster)


def test_large_chain_performance():
    # Large-scale test (within constraints): construct ~500-node linear chain.
    # Ensure the function returns the final node quickly and correctly.
    size = 500  # below 1000 per instructions
    nodes = [{"id": str(i)} for i in range(size)]
    # Create edges from 0->1, 1->2, ..., (size-2)->(size-1)
    edges = [{"source": str(i), "target": str(i + 1)} for i in range(size - 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 9.49ms -> 73.7μs (12770% faster)


def test_large_no_edges_returns_first_node():
    # Large-scale test with many nodes but no edges. Should return the first node.
    size = 800  # still below 1000
    nodes = [{"id": f"node_{i}"} for i in range(size)]
    edges: List[Dict[str, Any]] = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.75μs -> 709ns (288% 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-mkewfjzm and push.