⚡️ Speed up function find_leaf_nodes by 18,677%

[codeflash-ai]

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

⏱️ Runtime : 139 milliseconds → 742 microseconds (best of 503 runs)

📝 Explanation and details

Analysis of Profiling Results

• The overwhelming bottleneck is the double iteration: for each node, the code iterates over every edge, resulting in O(N * M) complexity (N = node count, M = edge count).
• The if edge["source"] == node["id"] check, with so many hits (millions), is the major time sink.

Optimization Strategy

• We can preprocess the set of all node ids that have an outgoing edge (i.e., any "source" in edges), and then simply filter nodes that are not in this set. This reduces the main logic to efficient O(N + M).

Installed Libraries

• No external libraries; only use built-in Python data structures for maximum compatibility.

---

Optimized Code

Explanation and Recommendations:

• Build a set of sources in O(M). Set lookup is O(1).
• Single pass over nodes in O(N) to select leaf nodes.
• This brings time complexity to O(N + M) instead of O(N*M).
• No external requirements or dependencies.

Summary:
This single code refactoring brings a dramatic speed-up to the function, removing the performance bottleneck shown in the profiler results (for edge in edges).
No compatibility concerns; this uses only core Python.
No change to function signature or return value.

---

If you need to handle extremely large datasets and memory is a constraint (very many edges), consider an iterator pattern, but for most practical cases the recommendation above is optimal.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

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():
    # One node, no edges: node is a leaf
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges) # 417ns -> 583ns (28.5% slower)

def test_two_nodes_one_edge():
    # Two nodes, one edge: only target is a leaf
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 750ns -> 708ns (5.93% faster)

def test_three_nodes_chain():
    # Three nodes in a chain: last node 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) # 958ns -> 792ns (21.0% faster)

def test_three_nodes_star():
    # One node points to two others: two 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 -> 750ns (33.3% faster)

def test_all_leaf_nodes():
    # No edges: all nodes are leaves
    nodes = [{"id": i} for i in range(5)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges) # 750ns -> 833ns (9.96% slower)

def test_no_leaf_nodes():
    # Every node has outgoing edge: no leaves
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 1}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 750ns -> 708ns (5.93% faster)

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

def test_empty_graph():
    # No nodes, no edges: no leaves
    nodes = []
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges) # 292ns -> 416ns (29.8% slower)

def test_edges_with_nonexistent_nodes():
    # Edges refer to nodes not in the node list: only listed nodes can be leaves
    nodes = [{"id": 1}]
    edges = [{"source": 2, "target": 1}, {"source": 3, "target": 4}]
    # Node 1 has no outgoing edges, so is a leaf
    codeflash_output = find_leaf_nodes(nodes, edges) # 625ns -> 708ns (11.7% slower)

def test_node_with_self_loop():
    # Node with edge to itself: not a leaf
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 1}]
    # Node 1 is not a leaf, node 2 is a leaf
    codeflash_output = find_leaf_nodes(nodes, edges) # 750ns -> 667ns (12.4% faster)

def test_duplicate_edges():
    # Multiple edges from one node to another: only targets are leaves
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 833ns -> 750ns (11.1% faster)

def test_isolated_node_with_edges_present():
    # Some nodes are isolated, others have edges
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 1, "target": 2}]
    # Node 1 is not a leaf, 2 and 3 are leaves (3 is isolated)
    codeflash_output = find_leaf_nodes(nodes, edges) # 917ns -> 750ns (22.3% faster)

def test_multiple_components():
    # Disconnected subgraphs
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}]
    edges = [{"source": 1, "target": 2}, {"source": 3, "target": 4}]
    # Leaves: 2 and 4
    codeflash_output = find_leaf_nodes(nodes, edges) # 1.12μs -> 833ns (35.1% faster)

def test_node_with_multiple_outgoing_edges():
    # Node with multiple outgoing edges is not 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) # 1.00μs -> 792ns (26.3% faster)

def test_nodes_with_non_integer_ids():
    # Node IDs are strings
    nodes = [{"id": "a"}, {"id": "b"}, {"id": "c"}]
    edges = [{"source": "a", "target": "b"}, {"source": "a", "target": "c"}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 1.29μs -> 750ns (72.3% faster)

def test_nodes_with_extra_data():
    # Nodes have extra fields, must preserve them
    nodes = [{"id": 1, "name": "A"}, {"id": 2, "name": "B"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 750ns -> 708ns (5.93% faster)

def test_edge_dict_with_extra_data():
    # Edges have extra fields, should be ignored
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2, "weight": 0.5}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 791ns -> 667ns (18.6% faster)

def test_nodes_with_none_id():
    # Node with id=None, should match edges with source=None
    nodes = [{"id": None}, {"id": 1}]
    edges = [{"source": None, "target": 1}]
    # Node with id=None has outgoing edge, so not a leaf; node 1 is a leaf
    codeflash_output = find_leaf_nodes(nodes, edges) # 916ns -> 708ns (29.4% faster)

def test_edges_with_none_source():
    # Edge with source=None, should not affect any node except id=None
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": None, "target": 1}]
    # Both nodes have no outgoing edges, so both are leaves
    codeflash_output = find_leaf_nodes(nodes, edges) # 833ns -> 708ns (17.7% faster)

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

def test_large_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) # 15.6ms -> 60.8μs (25618% faster)

def test_large_star():
    # 1 root, 999 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 # 30.5ms -> 57.1μs (53410% faster)

def test_large_fully_connected():
    # Every node has outgoing edges to all others: no leaves
    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]
    codeflash_output = find_leaf_nodes(nodes, edges) # 14.9ms -> 224μs (6547% faster)

def test_large_isolated_nodes():
    # 1000 isolated nodes, all are leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges) # 39.0μs -> 31.8μs (22.6% faster)

def test_large_sparse_graph():
    # 1000 nodes, 10 random edges, most nodes are leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    # Edges from 0 to 9 to 10 to 19
    edges = [{"source": i, "target": i+10} for i in range(10)]
    expected_leaves = [node for node in nodes if node["id"] not in range(10)]
    codeflash_output = find_leaf_nodes(nodes, edges) # 365μs -> 35.6μs (926% faster)

def test_large_multiple_components():
    # 10 chains of 100 nodes each
    chains = []
    edges = []
    for c in range(10):
        offset = c * 100
        chain = [{"id": offset + i} for i in range(100)]
        chains.extend(chain)
        edges.extend([{"source": offset + i, "target": offset + i + 1} for i in range(99)])
    # Each chain's last node is a leaf
    expected_leaves = [{"id": c * 100 + 99} for c in range(10)]
    codeflash_output = find_leaf_nodes(chains, edges) # 15.8ms -> 61.8μs (25451% 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 for find_leaf_nodes
# ------------------------

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

def test_single_node_no_edges():
    # One node, no edges: node is a leaf
    nodes = [{"id": 1}]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges) # 458ns -> 583ns (21.4% slower)

def test_two_nodes_one_edge():
    # Two nodes, one edge from 1 to 2: node 2 is a leaf
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 750ns -> 708ns (5.93% faster)

def test_three_nodes_chain():
    # 1 -> 2 -> 3: only node 3 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) # 1.00μs -> 791ns (26.4% faster)

def test_three_nodes_star():
    # 1 -> 2, 1 -> 3: nodes 2 and 3 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.04μs -> 792ns (31.6% 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) # 792ns -> 833ns (4.92% slower)

def test_no_nodes():
    # No nodes: should return empty list
    nodes = []
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges) # 250ns -> 417ns (40.0% slower)

def test_node_with_self_loop():
    # Node with a self-loop is not a leaf
    nodes = [{"id": 1}]
    edges = [{"source": 1, "target": 1}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 542ns -> 625ns (13.3% slower)

def test_node_with_incoming_but_no_outgoing():
    # Node 2 has only incoming edge, should be a leaf
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 750ns -> 708ns (5.93% faster)

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

def test_disconnected_graph():
    # Nodes not connected by any edge: all are leaves
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges) # 625ns -> 667ns (6.30% slower)

def test_cycle_graph():
    # 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) # 917ns -> 792ns (15.8% faster)

def test_multiple_edges_from_node():
    # 1 -> 2, 1 -> 3, 1 -> 4: 2,3,4 are leaves
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}]
    edges = [{"source": 1, "target": 2}, {"source": 1, "target": 3}, {"source": 1, "target": 4}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 1.25μs -> 875ns (42.9% 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) # 875ns -> 709ns (23.4% faster)

def test_isolated_and_connected_nodes():
    # 1 isolated, 2->3: 1 and 3 are leaves
    nodes = [{"id": 1}, {"id": 2}, {"id": 3}]
    edges = [{"source": 2, "target": 3}]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 958ns -> 792ns (21.0% faster)

def test_node_with_only_incoming_and_self_loop():
    # Node 2 has incoming from 1 and a self-loop: not a leaf
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 792ns -> 708ns (11.9% faster)

def test_nodes_with_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_leaf_nodes(nodes, edges) # 1.17μs -> 750ns (55.5% faster)

def test_edges_to_nonexistent_nodes():
    # Edges point to targets not in nodes: should not affect result
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}, {"source": 2, "target": 3}, {"source": 3, "target": 4}]
    # Only node 2 has no outgoing edge
    codeflash_output = find_leaf_nodes(nodes, edges) # 792ns -> 750ns (5.60% faster)

def test_nodes_with_extra_fields():
    # Nodes have extra fields, should be preserved in output
    nodes = [{"id": 1, "name": "a"}, {"id": 2, "name": "b"}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 750ns -> 708ns (5.93% faster)

def test_edges_with_extra_fields():
    # Edges have extra fields, should not affect result
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2, "weight": 5}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 708ns -> 667ns (6.15% faster)

def test_empty_edges_but_nodes_present():
    # All nodes are leaves if no edges
    nodes = [{"id": i} for i in range(10)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges) # 1.04μs -> 1.04μs (0.000% faster)

def test_large_ids():
    # Large integer ids
    nodes = [{"id": 10**12}, {"id": 10**12 + 1}]
    edges = [{"source": 10**12, "target": 10**12 + 1}]
    codeflash_output = find_leaf_nodes(nodes, edges) # 834ns -> 708ns (17.8% faster)

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

def test_large_linear_chain():
    # Large chain: only last node 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) # 15.7ms -> 61.0μs (25600% faster)

def test_large_star_graph():
    # One root, many leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = [{"source": 0, "target": i} for i in range(1, N)]
    leaves = [{"id": i} for i in range(1, N)]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 30.4ms -> 57.5μs (52790% faster)

def test_large_fully_disconnected():
    # All nodes are leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    edges = []
    codeflash_output = find_leaf_nodes(nodes, edges) # 38.8μs -> 31.8μs (22.0% faster)

def test_large_sparse_random_graph():
    # Random sparse edges, check correctness
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    random.seed(42)
    edges = []
    # Each node has a 1% chance of an outgoing edge to a random other node
    for i in range(N):
        if random.random() < 0.01:
            target = random.randint(0, N-1)
            if target != i:
                edges.append({"source": i, "target": target})
    # Find expected leaves: nodes with no outgoing edge
    outgoing = set(edge["source"] for edge in edges)
    expected = [node for node in nodes if node["id"] not in outgoing]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 322μs -> 38.1μs (748% faster)

def test_large_graph_with_isolated_and_connected_nodes():
    # Some nodes have no edges, some are internal, some are leaves
    N = 1000
    nodes = [{"id": i} for i in range(N)]
    # First 500 nodes: chain
    edges = [{"source": i, "target": i+1} for i in range(499)]
    # Next 250 nodes: star, node 500 as root
    edges += [{"source": 500, "target": i} for i in range(501, 751)]
    # Last 249 nodes: isolated
    # Leaves: node 499 (end of chain), nodes 501-750 (star leaves), nodes 751-999 (isolated)
    expected = [{"id": 499}] + [{"id": i} for i in range(501, 751)] + [{"id": i} for i in range(751, 1000)]
    codeflash_output = find_leaf_nodes(nodes, edges); result = codeflash_output # 15.6ms -> 53.5μs (29070% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from enum import EnumCheck
from src.dsa.nodes import find_leaf_nodes

def test_find_leaf_nodes():
    find_leaf_nodes([{'id': EnumCheck.UNIQUE}], [{'source': '\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00', 0: 0}])

def test_find_leaf_nodes_2():
    find_leaf_nodes([{'id': '', 'source': '\x00'}], [{'source': '\x00'}, {'source': ''}])

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