⚡️ Speed up function find_leaf_nodes by 18,209%

[codeflash-ai]

📄 18,209% (182.09x) speedup for find_leaf_nodes in src/dsa/nodes.py

⏱️ Runtime : 147 milliseconds → 804 microseconds (best of 133 runs)

📝 Explanation and details

Let's analyze and aggressively optimize your code based on the provided profiler.

Bottleneck

The main time cost comes from this loop.

This is O(N*M) where N is the number of nodes and M is the number of edges—very expensive.
The profiler confirms that nearly all of your program's time is spent inside the nested edge loop.

Faster Approach

Rather than scan all edges for each node, collect all node ids with outgoing edges first (i.e., all edge["source"]) using a set for O(1) lookups.

• Build a set of node ids with outgoing edges: source_ids = {edge["source"] for edge in edges}
• Any node whose "id" is not in source_ids is a leaf.

This reduces complexity from O(N*M) to O(N+M).

Code Rewrite

Notes

• No changes to the function signature or comments (existing relevant comment is kept).
• No custom libraries are required; compatible with the given environment.
• This implementation is maximal in speed and minimal in memory—just one intermediate set.

Summary of Optimization

• Reduced complexity from O(N*M) to O(N+M).
• Used highly optimized set lookups instead of repeated iteration.
• Output is identical to the original.

This rewrite should drastically reduce the runtime (by orders of magnitude for large inputs) compared to the original.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import random

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

# unit tests

# ----------------------
# 1. BASIC TEST CASES
# ----------------------

def test_single_node_no_edges():
    # Single node, no edges: node is a leaf
    nodes = [{"id": 1, "name": "A"}]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 500ns -> 666ns (24.9% slower)

def test_two_nodes_one_edge():
    # Two nodes, one edge: only target is a leaf
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 750ns -> 791ns (5.18% slower)

def test_three_nodes_chain():
    # Three nodes in a chain: only last is a leaf
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 3}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 959ns -> 875ns (9.60% faster)

def test_three_nodes_star():
    # One node points to all others: only targets are leaves
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.00μs -> 875ns (14.3% faster)

def test_all_nodes_are_leaves():
    # No edges: all nodes are leaves
    nodes = [{"id": i} for i in range(5)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 791ns -> 917ns (13.7% slower)

def test_no_leaves_all_point_somewhere():
    # Each node has outgoing edge: no leaves
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 3}, {"source": 3, "target": 1}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.04μs -> 916ns (13.8% faster)

def test_multiple_leaves():
    # Multiple leaves, non-consecutive ids
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.04μs -> 875ns (19.0% faster)

# ----------------------
# 2. EDGE TEST CASES
# ----------------------

def test_empty_nodes_and_edges():
    # Both lists empty: result empty
    nodes = []
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 333ns -> 458ns (27.3% slower)

def test_nodes_but_empty_edges():
    # Nodes present, no edges: all are leaves
    nodes = [{"id": i} for i in range(10)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.08μs -> 1.12μs (3.73% slower)

def test_edges_with_nonexistent_nodes():
    # Edges refer to nodes not in nodes list: should not affect output
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 3, "target": 4}, {"source": 3, "target": 1}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 875ns -> 875ns (0.000% faster)

def test_duplicate_node_ids():
    # Duplicate node ids: only the first occurrence is considered as per input order
    nodes = [{"id": 1}, {"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 875ns -> 750ns (16.7% faster)

def test_duplicate_edges():
    # Duplicate edges: should not affect result
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 833ns -> 791ns (5.31% faster)

def test_self_loop():
    # Node with self-loop is not a leaf
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 750ns -> 750ns (0.000% faster)

def test_disconnected_graph():
    # Some nodes not connected at all
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 917ns -> 833ns (10.1% faster)

def test_empty_edges_with_nonempty_nodes():
    # No edges, multiple nodes: all leaves
    nodes = [{"id": i} for i in range(3)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 584ns -> 709ns (17.6% slower)

def test_nodes_with_non_integer_ids():
    # Node ids are strings
    nodes = [{"id": "a"}, {"id": "b"}]
    edges = [{"source": "a", "target": "b"}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 916ns -> 791ns (15.8% faster)

def test_nodes_with_mixed_id_types():
    # Node ids are mixed types
    nodes = [{"id": 1}, {"id": "2"}, {"id": (3,)}]
    edges = [{"source": 1, "target": "2"}, {"source": "2", "target": (3,)}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.42μs -> 1.00μs (41.7% faster)

def test_edge_with_missing_source_or_target():
    # Edge missing 'source' or 'target' keys: should raise KeyError
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1}, {"target": 2}]
    with pytest.raises(KeyError):
        find_leaf_nodes(nodes, edges)

def test_nodes_with_extra_keys():
    # Nodes have extra irrelevant keys
    nodes = [{"id": 1, "foo": "bar"}, {"id": 2, "baz": 42}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 750ns -> 708ns (5.93% faster)

# ----------------------
# 3. LARGE SCALE TEST CASES
# ----------------------

def test_large_linear_chain():
    # 1000 nodes in a chain: only last is a leaf
    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_leaf_nodes(nodes, edges); result = codeflash_output # 16.3ms -> 65.7μs (24695% faster)

def test_large_star_graph():
    # 1 root node points to 999 others: all others are leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": 0, "target": i} for i in range(1, N)]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 31.0ms -> 60.7μs (50943% faster)

def test_large_fully_disconnected():
    # 1000 nodes, no edges: all leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 38.8μs -> 34.3μs (13.1% faster)

def test_large_cycle():
    # 1000 nodes in a cycle: no leaves
    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_leaf_nodes(nodes, edges); result = codeflash_output # 16.4ms -> 64.1μs (25463% faster)

def test_large_random_sparse_graph():
    # 1000 nodes, 500 random edges, check leaf count
    N = 1000
    E = 500
    nodes = [{"id": i} for i in range(N)]
    random.seed(42)
    edges = []
    used = set()
    while len(edges) < E:
        s = random.randint(0, N-1)
        t = random.randint(0, N-1)
        if s == t or (s, t) in used:
            continue
        edges.append({"source": s, "target": t})
        used.add((s, t))
    # Compute expected leaves
    sources = set(edge["source"] for edge in edges)
    expected_leaves = [node for node in nodes if node["id"] not in sources]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 14.6ms -> 76.3μs (19093% faster)

def test_large_graph_with_duplicate_edges():
    # 1000 nodes, 100 duplicate edges
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": 0, "target": i} for i in range(1, 101)] * 2  # duplicate edges
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 6.07ms -> 42.0μs (14347% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import random

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

# unit tests

# ------------------------------
# BASIC TEST CASES
# ------------------------------

def test_single_node_no_edges():
    # Single node, no edges: node is a leaf
    nodes = [{"id": 1, "name": "A"}]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 458ns -> 625ns (26.7% slower)

def test_two_nodes_one_edge():
    # Two nodes, one edge from A to B: only B is a leaf
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 792ns -> 750ns (5.60% faster)

def test_three_nodes_chain():
    # Three nodes in a chain: A->B->C; only C is a leaf
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "C"}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.21μs -> 917ns (31.8% faster)

def test_multiple_leaf_nodes():
    # Star topology: A->B, A->C, A->D; B, C, D are leaves
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}, {"id": "D"}]
    edges = [{"source": "A", "target": "B"}, {"source": "A", "target": "C"}, {"source": "A", "target": "D"}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.50μs -> 1.04μs (44.1% faster)

def test_all_nodes_are_leaves():
    # No edges, all nodes are leaves
    nodes = [{"id": i} for i in range(5)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 792ns -> 875ns (9.49% slower)

def test_no_leaves_all_nodes_have_outgoing():
    # Complete graph: every node has outgoing edge(s)
    nodes = [{"id": i} for i in range(3)]
    edges = [{"source": i, "target": j} for i in range(3) for j in range(3) if i != j]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.12μs -> 1.12μs (0.000% faster)

# ------------------------------
# EDGE TEST CASES
# ------------------------------

def test_empty_graph():
    # No nodes, no edges
    nodes = []
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 333ns -> 458ns (27.3% slower)

def test_edges_with_nonexistent_nodes():
    # Edges refer to node IDs not present in nodes
    nodes = [{"id": 1}]
    edges = [{"source": 2, "target": 1}, {"source": 1, "target": 3}]
    # Only node 1 is present, and it has an outgoing edge, so no leaves
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 667ns -> 750ns (11.1% slower)

def test_self_loop():
    # Node with a self-loop should not be a leaf
    nodes = [{"id": 1}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 541ns -> 708ns (23.6% slower)

def test_disconnected_nodes():
    # Some nodes are not connected at all
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}]
    # B and C have no outgoing edges, so both are leaves
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.00μs -> 833ns (20.0% faster)

def test_duplicate_edges():
    # Multiple identical edges, should not affect outcome
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 875ns -> 792ns (10.5% faster)

def test_nodes_with_extra_fields():
    # Nodes with extra fields should be preserved in output
    nodes = [{"id": 1, "label": "A", "data": 123}, {"id": 2, "label": "B", "data": 456}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 708ns -> 709ns (0.141% slower)

def test_edge_case_node_id_types():
    # Node IDs as strings, edges as integers (should not match)
    nodes = [{"id": "1"}, {"id": "2"}]
    edges = [{"source": 1, "target": 2}]
    # No node's id matches edge's source, so all nodes are leaves
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 833ns -> 750ns (11.1% faster)

def test_edge_case_duplicate_node_ids():
    # Duplicate node IDs in nodes list (should both be considered, but this is an invalid case)
    nodes = [{"id": 1}, {"id": 1}]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 583ns -> 667ns (12.6% slower)


def test_large_linear_chain():
    # 1000 nodes in a linear chain: 0->1->2->...->999
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": i+1} for i in range(N-1)]
    # Only the last node is a leaf
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 16.4ms -> 63.3μs (25853% faster)

def test_large_star_topology():
    # 1 center node, 999 leaves: 0->1, 0->2, ..., 0->999
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": 0, "target": i} for i in range(1, N)]
    # All nodes except 0 are leaves
    expected = [{"id": i} for i in range(1, N)]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 30.8ms -> 60.5μs (50921% faster)

def test_large_complete_graph():
    # 100 nodes, each node has outgoing edges to every other node
    N = 100
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": i, "target": j} for i in range(N) for j in range(N) if i != j]
    # No leaves
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 15.2ms -> 230μs (6481% faster)

def test_large_sparse_graph():
    # 1000 nodes, only 10 random edges
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    random.seed(42)
    edges = []
    for _ in range(10):
        src = random.randint(0, N-1)
        tgt = random.randint(0, N-1)
        while tgt == src:
            tgt = random.randint(0, N-1)
        edges.append({"source": src, "target": tgt})
    # All nodes except those with outgoing edges are leaves
    sources = set(edge["source"] for edge in edges)
    expected = [node for node in nodes if node["id"] not in sources]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 408μs -> 42.6μs (857% faster)

def test_large_all_leaves():
    # 1000 nodes, no edges: all nodes are leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 38.7μs -> 33.4μs (16.0% faster)

# ------------------------------
# ADDITIONAL EDGE CASES
# ------------------------------



def test_edge_case_multiple_targets():
    # Node with multiple outgoing edges, should not be a leaf
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.08μs -> 958ns (13.2% faster)

def test_edge_case_circular_graph():
    # Cycle: 1->2->3->1, no leaves
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 3}, {"source": 3, "target": 1}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.00μs -> 875ns (14.3% faster)

def test_edge_case_mixed_types():
    # Node IDs are mixed types, edges match only by exact type
    nodes = [{"id": 1}, {"id": "1"}]
    edges = [{"source": 1, "target": "1"}]
    # Node with id=1 has outgoing edge, node with id="1" does not
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 833ns -> 833ns (0.000% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from src.dsa.nodes import find_leaf_nodes
import pytest

def test_find_leaf_nodes():
    find_leaf_nodes([{'id': 3}], [{'source': 4, 0: ''}])

def test_find_leaf_nodes_2():
    with pytest.raises(KeyError):
        find_leaf_nodes([(v1 := {'source': 4, 'id': 4, '': 0}), v1, {}], [v1])

To edit these changes git checkout codeflash/optimize-find_leaf_nodes-mc8q8rz4 and push.