⚡️ Speed up function find_last_node by 21,941%

[codeflash-ai]

📄 21,941% (219.41x) speedup for find_last_node in src/algorithms/graph.py

⏱️ Runtime : 66.8 milliseconds → 303 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 220x speedup by eliminating a costly nested loop through a simple algorithmic improvement.

Key Optimization:

The original implementation uses all(e["source"] != n["id"] for e in edges) for each node, which creates an O(N × E) complexity where N is the number of nodes and E is the number of edges. For every node, it iterates through all edges to check if that node appears as a source.

The optimized version pre-computes a set of all source node IDs once: sources = {e["source"] for e in edges}. Then it uses O(1) set membership testing (n["id"] not in sources) for each node, reducing the overall complexity to O(N + E).

Performance Impact:

The speedup is most dramatic when the graph has many edges:

• Large linear chain (1000 nodes): 18.1ms → 54.6μs (331x faster)
• Large branching tree (511 nodes): 2.31ms → 18.2μs (126x faster)
• Large cyclic graph (500 nodes): 4.54ms → 27.7μs (163x faster)

For small graphs (2-4 nodes), the optimization still provides 50-80% speedup. Even the empty graph case shows only minimal overhead (14% slower), which is negligible given the microsecond timescale.

Why It's Faster:

Set construction and lookup in Python are highly optimized hash table operations with O(1) average-case complexity. The original approach repeatedly scans the entire edges list, which becomes prohibitively expensive as the graph grows. The optimization trades a small upfront cost (building the set) for massive savings during the node iteration.

This optimization is universally beneficial across all test cases involving non-trivial edge counts, making it an excellent candidate for merging.

✅ 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

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)  # 1.21μs -> 958ns (26.1% 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)  # 1.71μs -> 1.12μs (51.8% faster)


def test_three_nodes_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)  # 2.17μs -> 1.21μs (79.3% faster)


def test_three_nodes_branch():
    # 1 -> 2, 1 -> 3: last nodes are 2 and 3 (should return 2, as per implementation)
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}]
    # Should return the first node (2) with no outgoing edges
    codeflash_output = find_last_node(nodes, edges)  # 1.83μs -> 1.21μs (51.7% faster)


def test_no_nodes():
    # No nodes at all: should return None
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 750ns -> 875ns (14.3% slower)


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


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


def test_multiple_last_nodes():
    # Multiple nodes with no outgoing edges: returns the first one found
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}]
    # Both B and C have no outgoing edges; function returns B (first in list)
    codeflash_output = find_last_node(nodes, edges)  # 1.92μs -> 1.08μs (76.9% 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)  # 1.83μs -> 1.21μs (51.7% faster)


def test_edges_with_nonexistent_nodes():
    # Edges refer to nodes not in the nodes list: should not affect result
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 3}, {"source": 3, "target": 2}]
    # Node 2 has no outgoing edges, so should return node 2
    codeflash_output = find_last_node(nodes, edges)  # 1.83μs -> 1.17μs (57.2% faster)


def test_duplicate_node_ids():
    # Duplicate node IDs: function should return the first one with no outgoing edges
    nodes = [{"id": 1}, {"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    # Both nodes with id=1 have outgoing edges, node 2 does not
    codeflash_output = find_last_node(nodes, edges)  # 2.00μs -> 1.12μs (77.8% faster)


def test_nodes_with_extra_fields():
    # Nodes have extra fields, should still work
    nodes = [{"id": 1, "value": "x"}, {"id": 2, "value": "y"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)  # 1.71μs -> 1.08μs (57.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)  # 1.75μs -> 1.08μs (61.6% faster)


def test_node_id_is_string():
    # Node IDs are strings
    nodes = [{"id": "foo"}, {"id": "bar"}]
    edges = [{"source": "foo", "target": "bar"}]
    codeflash_output = find_last_node(nodes, edges)  # 1.83μs -> 1.12μs (62.9% faster)


def test_node_id_is_tuple():
    # Node IDs are tuples
    nodes = [{"id": (1, 2)}, {"id": (2, 3)}]
    edges = [{"source": (1, 2), "target": (2, 3)}]
    codeflash_output = find_last_node(nodes, edges)  # 1.88μs -> 1.17μs (60.8% 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(N)]
    edges = [{"source": i, "target": i + 1} for i in range(N - 1)]
    codeflash_output = find_last_node(nodes, edges)  # 18.1ms -> 54.6μs (33085% faster)


def test_large_branching_tree():
    # Large binary tree: all leaves are last nodes, should return the first leaf
    N = 511  # Full binary tree with 255 inner nodes and 256 leaves
    nodes = [{"id": i} for i in range(N)]
    edges = []
    for i in range((N - 1) // 2):
        edges.append({"source": i, "target": 2 * i + 1})
        edges.append({"source": i, "target": 2 * i + 2})
    # Leaves are nodes with id >= (N-1)//2 + 1
    first_leaf_id = (N - 1) // 2 + 1
    codeflash_output = find_last_node(nodes, edges)  # 2.31ms -> 18.2μs (12610% faster)


def test_large_graph_with_no_last_node():
    # Large cyclic graph: all nodes have outgoing edges
    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)  # 4.54ms -> 27.7μs (16270% faster)


def test_large_graph_with_extra_fields():
    # Large graph, nodes and edges have extra fields
    N = 200
    nodes = [{"id": i, "meta": f"node_{i}"} for i in range(N)]
    edges = [{"source": i, "target": i + 1, "weight": i * 2} for i in range(N - 1)]
    codeflash_output = find_last_node(nodes, edges)  # 748μs -> 10.8μs (6840% faster)


def test_large_disconnected_nodes():
    # Many nodes, only a few connected, most are isolated (should return first isolated node)
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": 0, "target": 1}]
    # All nodes except 0 have no outgoing edges, so first is node 1
    codeflash_output = find_last_node(nodes, edges)  # 1.88μs -> 1.21μs (55.2% 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 itself
    nodes = [{"id": 1, "name": "A"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 1.21μs -> 959ns (26.0% faster)


def test_two_nodes_one_edge():
    # Two nodes, one edge from A to B: should return B (no outgoing edges)
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges)  # 1.83μs -> 1.17μs (57.3% faster)


def test_three_nodes_linear_chain():
    # Three nodes in a linear chain: A -> B -> C, should return C
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "C"}]
    codeflash_output = find_last_node(nodes, edges)  # 2.21μs -> 1.25μs (76.6% faster)


def test_multiple_nodes_multiple_edges():
    # Diamond shape: A -> B, A -> C, B -> D, C -> D; last node is D
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}, {"id": "D"}]
    edges = [
        {"source": "A", "target": "B"},
        {"source": "A", "target": "C"},
        {"source": "B", "target": "D"},
        {"source": "C", "target": "D"},
    ]
    codeflash_output = find_last_node(nodes, edges)  # 2.92μs -> 1.29μs (126% faster)


def test_multiple_last_nodes_returns_first():
    # Two nodes with no outgoing edges: should return the first in nodes order
    nodes = [{"id": 1}, {"id": 2}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 1.21μs -> 959ns (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)  # 708ns -> 834ns (15.1% slower)


def test_nodes_but_no_edges():
    # Multiple nodes, no edges: should return the first node
    nodes = [{"id": "X"}, {"id": "Y"}, {"id": "Z"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 1.21μs -> 959ns (26.1% 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)  # 1.88μs -> 1.25μs (50.0% faster)


def test_circular_graph():
    # Circular reference: A -> B, B -> C, C -> A; should return None
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [
        {"source": "A", "target": "B"},
        {"source": "B", "target": "C"},
        {"source": "C", "target": "A"},
    ]
    codeflash_output = find_last_node(nodes, edges)  # 2.33μs -> 1.29μs (80.7% faster)


def test_disconnected_graph():
    # Disconnected: A->B, C (no edges): should return C (no outgoing edges)
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges)  # 1.88μs -> 1.12μs (66.7% faster)


def test_multiple_nodes_with_no_outgoing_edges():
    # Multiple nodes with no outgoing edges, returns first in nodes order
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges)  # 1.21μs -> 1.00μs (20.9% faster)


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


def test_node_with_multiple_outgoing_edges():
    # Node with multiple outgoing edges, only one node with none
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}, {"source": "A", "target": "C"}]
    codeflash_output = find_last_node(nodes, edges)  # 1.92μs -> 1.21μs (58.5% faster)


def test_edges_with_nonexistent_nodes():
    # Edges refer to nodes not in nodes list; should ignore those edges
    nodes = [{"id": "X"}]
    edges = [{"source": "Y", "target": "Z"}]
    codeflash_output = find_last_node(nodes, edges)  # 1.42μs -> 1.08μs (30.8% faster)


def test_nodes_with_additional_properties():
    # Nodes have extra properties, should return full node dict
    nodes = [{"id": 1, "label": "foo"}, {"id": 2, "label": "bar"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges)  # 1.83μs -> 1.12μs (62.9% faster)


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


def test_large_linear_chain():
    # Large chain: 1000 nodes, linear edges; last node is the last in the list
    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)  # 18.2ms -> 54.5μs (33226% faster)


def test_large_star_graph():
    # 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)]
    # Should return the first node with no outgoing edges, which is 1
    codeflash_output = find_last_node(nodes, edges)  # 37.2μs -> 19.6μs (90.2% faster)


def test_large_graph_all_nodes_have_outgoing():
    # Every node has at least one outgoing edge, 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)  # 18.3ms -> 54.2μs (33583% faster)


def test_large_graph_multiple_last_nodes():
    # Half nodes have outgoing, half do not; returns first last node
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i + 1} for i in range(N // 2)]
    # First node with no outgoing edge is N//2
    codeflash_output = find_last_node(nodes, edges)  # 4.57ms -> 27.7μs (16418% faster)


def test_large_graph_with_disconnected_components():
    # Two chains, one longer than the other; should return the last node in the first chain
    nodes = [{"id": i} for i in range(10)] + [{"id": 100 + i} for i in range(20)]
    edges = [{"source": i, "target": i + 1} for i in range(9)] + [
        {"source": 100 + i, "target": 100 + i + 1} for i in range(19)
    ]
    # Should return {'id': 9} (last node of first chain, as it appears first in nodes)
    codeflash_output = find_last_node(nodes, edges)  # 6.50μs -> 2.58μs (152% 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-mjj7rfyx and push.