⚡️ Speed up function find_last_node by 15,521%

[codeflash-ai]

📄 15,521% (155.21x) speedup for find_last_node in src/dsa/nodes.py

⏱️ Runtime : 115 milliseconds → 735 microseconds (best of 591 runs)

📝 Explanation and details

Here's an optimized version of your code.
Your profiling result shows the code is dominated by the repeated check of whether any edge references a node’s id as "source".
This is O(N*M) where N = nodes, M = edges.
We can precompute the set of all edge "source" ids so the inner check is O(1) instead of O(M).

Here’s the faster rewrite (comments preserved).

Summary of speedup.

• The inner check is O(1) per node, not O(#edges)
• The total runtime is O(N + M), not O(N*M)
• Memory: proportional to #edges (edge_sources set)

Behavior and return values are unchanged.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from src.dsa.nodes 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); result = codeflash_output # 1.62μs -> 1.38μs (18.2% faster)

def test_two_nodes_one_edge():
    # Two nodes, one edge from node 1 to node 2: node 2 is last (no outgoing edges)
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.46μs -> 1.54μs (59.4% faster)

def test_three_nodes_linear_chain():
    # 1 -> 2 -> 3, node 3 is last
    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 # 3.12μs -> 1.67μs (87.5% faster)

def test_multiple_last_nodes_returns_first():
    # Two nodes with no outgoing edges: should return the first such node
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    # Both node 2 and 3 have no outgoing edges, so node 2 is returned (first in list)
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.25μs -> 1.50μs (50.0% faster)

def test_no_nodes():
    # No nodes at all: should return None
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 709ns -> 792ns (10.5% slower)

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

def test_no_edges_multiple_nodes():
    # Multiple nodes, no edges: should return the first node
    nodes = [{"id": "x"}, {"id": "y"}, {"id": "z"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 1.58μs -> 1.38μs (15.1% faster)

def test_all_nodes_with_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); result = codeflash_output # 3.08μs -> 1.25μs (147% faster)

def test_cycle_graph():
    # Nodes form a cycle: no last node
    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); result = codeflash_output # 3.00μs -> 1.21μs (148% faster)

def test_node_with_self_loop():
    # Node with a self-loop: should not be last node
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 1}]
    # Node 2 has no outgoing edges, so it is the last node
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.38μs -> 1.58μs (50.0% faster)

def test_edges_with_nonexistent_nodes():
    # Edge refers to a node not in the nodes list: should ignore such edges
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 3}]  # node 3 does not exist
    # Node 2 has no outgoing edges, so it is the last node
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.38μs -> 1.58μs (50.0% faster)

def test_nodes_with_non_integer_ids():
    # Node IDs are strings, not integers
    nodes = [{"id": "foo"}, {"id": "bar"}]
    edges = [{"source": "foo", "target": "bar"}]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.33μs -> 1.58μs (47.4% faster)

def test_duplicate_node_ids():
    # Duplicate node IDs: should return the first last-node found
    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); result = codeflash_output # 2.88μs -> 1.62μs (76.9% faster)

def test_empty_edges_with_empty_nodes():
    # Both nodes and edges are empty
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 708ns -> 792ns (10.6% slower)

def test_edges_empty_dicts():
    # Edges are empty dicts: should not match any source, so all nodes are last nodes, return first
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{}]
    # Since the edge dict does not have "source", all(e["source"] != n["id"] for e in edges) will raise KeyError
    # So this is an edge case for robustness: should raise KeyError
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)

def test_nodes_with_extra_fields():
    # Nodes have extra fields; function should ignore them
    nodes = [{"id": 1, "name": "A", "data": [1,2,3]}, {"id": 2, "foo": "bar"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.42μs -> 1.58μs (52.6% faster)

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

def test_large_linear_chain():
    # 1000 nodes in a linear chain: last node should be 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); result = codeflash_output # 24.9ms -> 62.3μs (39830% faster)

def test_large_branching_graph():
    # 500 nodes, each with two outgoing edges except the last 2 nodes
    N = 500
    nodes = [{"id": i} for i in range(N)]
    edges = []
    for i in range(N-2):
        edges.append({"source": i, "target": i+1})
        edges.append({"source": i, "target": i+2})
    # Last two nodes have no outgoing edges, so the first of those is returned
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output

def test_large_graph_no_last_node():
    # 1000 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 # 25.1ms -> 61.0μs (40978% faster)

def test_large_sparse_graph():
    # 1000 nodes, only 10 edges, so most nodes are last nodes; should return the first node with no outgoing edge
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i+1} for i in range(10)]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 9.67μs -> 2.33μs (314% faster)

def test_large_all_nodes_no_edges():
    # 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.67μs -> 1.46μs (14.3% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
from src.dsa.nodes import find_last_node

# unit tests

# ---------------------------
# Basic Test Cases
# ---------------------------

def test_single_node_no_edges():
    # One node, no edges: node should be returned as last node
    nodes = [{"id": "A"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 1.54μs -> 1.38μs (12.1% faster)

def test_two_nodes_one_edge():
    # Two nodes, one edge from A to B: B is the last node
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.33μs -> 1.54μs (51.4% faster)

def test_three_nodes_linear_chain():
    # Three nodes in a chain: A -> B -> C, C is last node
    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 # 3.12μs -> 1.67μs (87.5% faster)

def test_branching_graph():
    # Branching: A -> B, A -> C, B -> D, C -> D, D is last node
    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); result = codeflash_output # 3.88μs -> 1.75μs (121% faster)

def test_multiple_terminal_nodes_returns_first():
    # Two terminal nodes (no outgoing edges): returns the first in nodes
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}]  # C is disconnected, B is terminal
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.38μs -> 1.54μs (54.0% 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 # 708ns -> 792ns (10.6% slower)

def test_nodes_with_no_matching_edges():
    # Nodes present, edges refer to non-existent sources
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "X", "target": "A"}]  # "X" not in nodes
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 1.79μs -> 1.50μs (19.5% faster)

def test_all_nodes_with_outgoing_edges():
    # Every node has at least one outgoing edge: no terminal node
    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.42μs -> 1.12μs (115% faster)

def test_cycle_graph():
    # Cycle: A -> B -> C -> A, no terminal node
    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); result = codeflash_output # 3.08μs -> 1.25μs (147% faster)

def test_node_with_self_loop():
    # Node with self-loop is not terminal
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [
        {"source": "A", "target": "A"}
    ]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.25μs -> 1.54μs (46.0% faster)

def test_disconnected_nodes():
    # Nodes not connected by any edges
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 1.58μs -> 1.38μs (15.1% faster)

def test_multiple_edges_from_one_node():
    # One node with multiple outgoing edges, only one terminal
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [
        {"source": "A", "target": "B"},
        {"source": "A", "target": "C"}
    ]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.42μs -> 1.54μs (56.7% faster)

def test_edge_with_missing_source_key():
    # Edge missing 'source' key should raise KeyError
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"target": "B"}]
    with pytest.raises(KeyError):
        find_last_node(nodes, edges)

def test_node_with_additional_fields():
    # Nodes can have extra fields, should not affect result
    nodes = [{"id": "A", "value": 1}, {"id": "B", "value": 2}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.46μs -> 1.58μs (55.3% faster)

def test_edge_with_additional_fields():
    # Edges can have extra fields, should not affect result
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B", "weight": 10}]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 2.42μs -> 1.58μs (52.7% faster)

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

def test_large_linear_chain():
    # Large chain: 1000 nodes, each points to next, last node is terminal
    n = 1000
    nodes = [{"id": str(i)} for i in range(n)]
    edges = [{"source": str(i), "target": str(i+1)} for i in range(n-1)]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 26.6ms -> 110μs (24016% faster)

def test_large_star_topology():
    # Star: node 0 points to all others, all others are terminal
    n = 1000
    nodes = [{"id": str(i)} for i in range(n)]
    edges = [{"source": "0", "target": str(i)} for i in range(1, n)]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 56.7μs -> 23.6μs (140% faster)

def test_large_fully_connected_graph():
    # Fully connected: every node has outgoing edges to every other node (no terminal)
    n = 100
    nodes = [{"id": str(i)} for i in range(n)]
    edges = [{"source": str(i), "target": str(j)} for i in range(n) for j in range(n) if i != j]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 25.5ms -> 377μs (6645% faster)

def test_large_disconnected_graph():
    # All nodes have no edges, all are terminal, first node should be returned
    n = 500
    nodes = [{"id": str(i)} for i in range(n)]
    edges = []
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 1.71μs -> 1.50μs (13.9% faster)

def test_large_graph_with_multiple_terminal_nodes():
    # 100 nodes, first 90 each point to next, last 10 are disconnected
    n = 100
    nodes = [{"id": str(i)} for i in range(n)]
    edges = [{"source": str(i), "target": str(i+1)} for i in range(89)]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 259μs -> 11.3μs (2195% 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-mce2fvtu and push.