⚡️ Speed up function find_last_node by 25,446%

[codeflash-ai]

📄 25,446% (254.46x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 92.3 milliseconds → 361 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 254x speedup by eliminating a nested loop complexity issue. Here's why:

The Core Problem:
The original implementation uses a nested comprehension that checks all(e["source"] != n["id"] for e in edges) for each node. This creates O(N × M) comparisons where N is the number of nodes and M is the number of edges. For every node candidate, the code must scan through all edges repeatedly.

The Optimization:
The optimized version pre-computes a set of source IDs: source_ids = {e["source"] for e in edges}. This transforms the problem into:

1. One-time O(M) operation to build the set
2. O(N) lookups with O(1) average-case set membership checks

This reduces the overall complexity from O(N × M) to O(N + M).

Why This Matters:

• Set membership (in) vs repeated iteration: Python sets use hash tables, making lookups nearly instantaneous compared to iterating through a list for each check.
• Single pass through edges: Building the set once is far cheaper than the original code's repeated iteration through all edges for every node.

Performance Impact by Test Case:

• Small graphs (2-10 nodes/edges): 30-93% faster - modest gains as overhead is minimal
• Medium graphs (100 nodes/edges): 3000%+ faster - the O(N×M) penalty becomes significant
• Large graphs (1000 nodes/edges): 32,000%+ faster - the nested loop becomes catastrophic. For example, test_large_linear_chain drops from 18.6ms to 57.2μs because it eliminates ~1 million comparisons.

Special Cases:

• Empty graphs see slight slowdown (13-29%) due to set creation overhead when there's nothing to optimize
• Graphs with cycles or multiple sinks benefit equally since the improvement is in the lookup mechanism, not the logic

The optimization is universally beneficial for any non-trivial graph workload and essential for production code processing moderate to large graphs.

✅ 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

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 as the last node
    nodes = [{"id": 1, "label": "A"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.29μs -> 958ns (34.8% faster)


def test_two_nodes_one_edge():
    # Two nodes, one edge from 1 -> 2: should return node 2 as the last 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.79μs -> 1.17μs (53.7% faster)


def test_three_nodes_linear_chain():
    # Three nodes, edges: 1->2, 2->3. Last node should be 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.21μs (82.8% faster)


def test_multiple_sinks_returns_first():
    # Multiple nodes with no outgoing edges: returns the first one found
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    # Both node 2 and 3 have no outgoing edges, should return node 2 (first found)
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.75μs -> 1.12μs (55.6% 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 -> 917ns (13.7% slower)


def test_nodes_no_edges():
    # Multiple nodes, no edges: should return the first node
    nodes = [{"id": 10}, {"id": 20}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 958ns (30.5% faster)


def test_cycle_graph():
    # Cycle: 1->2, 2->3, 3->1. No node is a sink, 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.3% faster)


def test_disconnected_graph():
    # Disconnected graph: one isolated node (no edges)
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    # node 3 is isolated, node 2 is a sink, both have no outgoing edges, should return node 2
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.79μs -> 1.12μs (59.2% faster)


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


def test_node_with_multiple_outgoing_edges():
    # Node with multiple outgoing edges, only one sink
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}]
    # nodes 2 and 3 have no outgoing edges, should return node 2 (first found)
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.21μs (55.2% faster)


def test_non_integer_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.50μs -> 1.29μs (93.6% faster)


def test_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)  # 2.08μs -> 875ns (138% faster)


def test_missing_target_key():
    # Edge missing 'target' key should not affect function (since only 'source' is checked)
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1}]
    # node 2 has no outgoing edges, should return node 2
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.88μs -> 1.12μs (66.7% faster)


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


def test_large_linear_chain():
    # Large linear chain: 1->2->3->...->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.6ms -> 57.2μs (32322% faster)


def test_large_star_graph():
    # Star graph: node 0 points to all others, all others are sinks
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": 0, "target": i} for i in range(1, N)]
    # First sink node is node 1
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 39.3μs -> 20.9μs (88.4% faster)


def test_large_disconnected_nodes():
    # 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.2% faster)


def test_large_no_sink():
    # 1000 nodes in a cycle: no sinks, 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  # 18.3ms -> 56.0μs (32642% faster)


def test_large_multiple_sinks():
    # 1000 nodes, first 990 are sources, last 10 are sinks
    N = 1000
    K = 10
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i + 1} for i in range(N - K)]
    # The first sink is node N-K+1
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 18.1ms -> 55.7μs (32464% 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: should return the node itself
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.25μs -> 959ns (30.3% faster)


def test_two_nodes_one_edge():
    # Two nodes, one edge from 1->2: last node is node 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.12μs (62.9% 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.21μs -> 1.21μs (82.6% faster)


def test_multiple_end_nodes():
    # Four nodes, edges: 1->2, 1->3. Both 2 and 3 are possible last nodes (no outgoing edges).
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.83μs -> 1.21μs (51.8% faster)


def test_isolated_node_among_others():
    # Nodes 1->2, node 3 is isolated. Should return 2 or 3.
    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.08μs (61.4% faster)


# 2. 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  # 709ns -> 834ns (15.0% slower)


def test_edges_but_no_nodes():
    # Edges present, but no nodes: should return None
    nodes = []
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 709ns -> 1.00μs (29.1% slower)


def test_all_nodes_have_outgoing_edges():
    # Every node has an outgoing edge (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.21μs -> 1.29μs (70.9% faster)


def test_duplicate_edges():
    # Duplicate edges should not affect the result
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [
        {"source": 1, "target": 2},
        {"source": 1, "target": 2},
        {"source": 2, "target": 3},
    ]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 2.25μs -> 1.33μs (68.8% faster)


def test_node_with_self_loop():
    # Node with self-loop is not a 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.3% faster)


def test_multiple_isolated_nodes():
    # All nodes are isolated, any node can be returned
    nodes = [{"id": "a"}, {"id": "b"}, {"id": "c"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.21μs -> 1.04μs (16.0% faster)


def test_non_integer_node_ids():
    # Node ids are strings
    nodes = [{"id": "start"}, {"id": "end"}]
    edges = [{"source": "start", "target": "end"}]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.96μs -> 1.21μs (62.1% faster)


def test_extra_node_fields():
    # Nodes can have extra fields, but match by 'id'
    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.75μs -> 1.12μs (55.6% faster)


# 3. Large Scale Test Cases


def test_large_linear_chain():
    # Large chain of nodes: 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.5ms -> 56.9μs (32363% faster)


def test_large_star_graph():
    # Node 0 points to all others, all others have no outgoing edges
    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  # 39.2μs -> 20.7μs (89.5% faster)


def test_large_all_isolated_nodes():
    # All nodes are isolated, any can be returned
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = []
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 1.38μs -> 1.08μs (26.8% faster)


def test_large_cycle():
    # All nodes in a cycle (no last node)
    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  # 18.5ms -> 56.3μs (32746% faster)


def test_performance_many_edges():
    # Dense graph, but one node is a sink
    N = 100
    nodes = [{"id": i} for i in range(N)]
    # All nodes except last point to last node
    edges = [{"source": i, "target": N - 1} for i in range(N - 1)]
    codeflash_output = find_last_node(nodes, edges)
    result = codeflash_output  # 207μs -> 6.62μs (3031% 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-mjhql0kq and push.