Skip to content

⚡️ Speed up function find_last_node by 6% #20

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed

Conversation

codeflash-ai[bot]
Copy link

@codeflash-ai codeflash-ai bot commented Jun 21, 2025

📄 6% (0.06x) speedup for find_last_node in src/dsa/nodes.py

⏱️ Runtime : 790 microseconds 743 microseconds (best of 328 runs)

📝 Explanation and details

Here is an optimized version of your program.

Optimization rationale

  • The set comprehension for source_ids = {e["source"] for e in edges} is already a good approach.
  • However, calling n["id"] not in source_ids for every node can be costly, especially if nodes is large.
  • The best way to speed this up is to turn nodes into an iterable of node IDs for efficient lookup/comparison, but functionally, the code is nearly minimal.
  • To slightly speed up.
    • Pre-fetch field lookups to local variables if inside a comprehension.
    • Avoid dict lookup in tight loop by extracting n["id"] once.
    • Combine all into a single fast loop without generator expressions.

That said, the next((...), None) daily pattern is already efficient. But you can squeeze out a bit more performance by:

  • Changing the node scan to scan nodes only once but precompute the set before.
  • Use set subtraction between node IDs and source IDs, and then find the node object for the difference.

Below is the code.

Why is this faster?

  • No per-node "id" lookup on every generator iteration.
  • Only one pass through nodes, one pass through edges.
  • Pure set operations and dict lookups are the fastest way in Python for these container types.

If you are guaranteed there is only ever one possible "last" node, you could even do.

But the first (set) version will slightly outperform your original for large lists, as it avoids repeating lookups.

Choose the version that best suits your typical data size and cardinality.

Correctness verification report:

Test Status
⚙️ Existing Unit Tests 🔘 None Found
🌀 Generated Regression Tests 45 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():
    # Single node, no edges: should return the node itself
    nodes = [{"id": "A"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges) # 1.42μs -> 500ns (183% faster)

def test_two_nodes_one_edge():
    # Two nodes, one edge from A to B: should return B
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 667ns (131% faster)

def test_three_nodes_linear_chain():
    # Linear chain: A -> B -> C; last node is C
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "C"}]
    codeflash_output = find_last_node(nodes, edges) # 1.67μs -> 791ns (111% faster)

def test_multiple_leaves_returns_first_leaf():
    # Tree: A -> B, A -> C; both B and C are leaves, should return B (first in nodes)
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}, {"source": "A", "target": "C"}]
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 750ns (106% faster)

def test_cycle_graph():
    # Cycle: A -> B -> C -> A; all nodes are sources, so no last node, 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) # 1.29μs -> 833ns (55.0% faster)

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

def test_empty_nodes_and_edges():
    # No nodes or edges: should return None
    nodes = []
    edges = []
    codeflash_output = find_last_node(nodes, edges) # 875ns -> 375ns (133% faster)

def test_nodes_with_no_outgoing_edges():
    # All nodes have no outgoing edges, so all are leaves; should return first node
    nodes = [{"id": "X"}, {"id": "Y"}, {"id": "Z"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges) # 1.42μs -> 500ns (183% faster)

def test_edges_with_nonexistent_source():
    # Edge source not in nodes: should not affect result, all nodes are leaves
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "C", "target": "A"}]
    codeflash_output = find_last_node(nodes, edges) # 1.46μs -> 625ns (133% faster)

def test_edges_with_nonexistent_target():
    # Edge target not in nodes: should not affect result, B is not a source
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "C"}]
    codeflash_output = find_last_node(nodes, edges) # 1.50μs -> 708ns (112% faster)

def test_duplicate_edges():
    # Duplicate edges should not affect result
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [
        {"source": "A", "target": "B"},
        {"source": "A", "target": "B"}
    ]
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 708ns (118% faster)

def test_self_loop():
    # Self-loop: node points to itself; should not be a leaf
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "A"}]
    codeflash_output = find_last_node(nodes, edges) # 1.50μs -> 708ns (112% faster)

def test_all_nodes_are_sources():
    # Every node is a source: no leaves, should return None
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [
        {"source": "A", "target": "B"},
        {"source": "B", "target": "A"}
    ]
    codeflash_output = find_last_node(nodes, edges) # 1.17μs -> 708ns (64.8% faster)

def test_node_with_multiple_incoming_edges():
    # Node C has two incoming edges, should still be last node
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [
        {"source": "A", "target": "C"},
        {"source": "B", "target": "C"}
    ]
    codeflash_output = find_last_node(nodes, edges) # 1.58μs -> 750ns (111% faster)

def test_node_with_multiple_outgoing_edges():
    # Node A has two outgoing edges, B and C are both leaves, should return B
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [
        {"source": "A", "target": "B"},
        {"source": "A", "target": "C"}
    ]
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 708ns (118% faster)

def test_node_id_is_integer():
    # Node IDs as integers instead of strings
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 750ns (106% faster)

def test_node_id_is_tuple():
    # Node IDs as tuples
    nodes = [{"id": (1, "A")}, {"id": (2, "B")}]
    edges = [{"source": (1, "A"), "target": (2, "B")}]
    codeflash_output = find_last_node(nodes, edges) # 1.67μs -> 833ns (100% faster)

def test_nodes_with_extra_attributes():
    # Nodes have extra attributes, should return full dict
    nodes = [{"id": "A", "value": 1}, {"id": "B", "value": 2}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges) # 1.50μs -> 708ns (112% faster)

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

def test_large_linear_chain():
    # Large linear chain: 1000 nodes, last node should be the last in the chain
    nodes = [{"id": str(i)} for i in range(1000)]
    edges = [{"source": str(i), "target": str(i + 1)} for i in range(999)]
    codeflash_output = find_last_node(nodes, edges) # 112μs -> 109μs (2.16% faster)

def test_large_tree_multiple_leaves():
    # Large tree: root with 999 leaf nodes
    nodes = [{"id": "root"}] + [{"id": f"leaf{i}"} for i in range(999)]
    edges = [{"source": "root", "target": f"leaf{i}"} for i in range(999)]
    # All leaves are leaves; should return the first leaf node
    codeflash_output = find_last_node(nodes, edges) # 25.2μs -> 23.5μs (7.09% faster)

def test_large_graph_no_leaves():
    # 1000 nodes in a cycle: no leaves, should return None
    nodes = [{"id": str(i)} for i in range(1000)]
    edges = [{"source": str(i), "target": str((i + 1) % 1000)} for i in range(1000)]
    codeflash_output = find_last_node(nodes, edges) # 111μs -> 108μs (2.57% faster)

def test_large_graph_with_extra_attributes():
    # 1000 nodes, each with extra attributes, linear chain
    nodes = [{"id": str(i), "data": i*2} for i in range(1000)]
    edges = [{"source": str(i), "target": str(i + 1)} for i in range(999)]
    codeflash_output = find_last_node(nodes, edges) # 110μs -> 108μs (2.04% faster)

# -------------------- Mutation Testing Traps --------------------

def test_returns_first_leaf_even_if_not_last_in_edges():
    # Should return first leaf node in nodes, not in edges
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "C"}, {"source": "A", "target": "B"}]
    # Both B and C are leaves, but B comes first in nodes
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 708ns (118% faster)

def test_none_returned_if_no_leaf():
    # All nodes are sources, no leaves, 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) # 1.25μs -> 792ns (57.8% faster)


def test_ignores_edges_with_none_source():
    # Edges with source=None should not affect leaf detection
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": None, "target": "A"}]
    codeflash_output = find_last_node(nodes, edges) # 1.62μs -> 667ns (144% faster)

def test_nodes_with_duplicate_ids():
    # If nodes have duplicate ids, should return the first one not in source_ids
    nodes = [{"id": "A"}, {"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges) # 1.62μs -> 709ns (129% 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: the only node is the last node
    nodes = [{"id": "A"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges) # 1.38μs -> 541ns (154% faster)

def test_two_nodes_one_edge():
    # Two nodes, one edge from A to B: B is the last node (no outgoing edges)
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges) # 1.46μs -> 708ns (106% faster)

def test_three_nodes_linear():
    # Linear graph: A -> B -> C; C is the 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) # 1.58μs -> 750ns (111% faster)

def test_three_nodes_branch():
    # Branch: A -> B, A -> C; B and C are both last nodes, function returns first found
    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 # 1.54μs -> 750ns (105% 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) # 875ns -> 375ns (133% faster)

def test_nodes_no_edges():
    # Multiple nodes, no edges: any node is a last node (no outgoing edges)
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = []
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 1.38μs -> 541ns (154% faster)

def test_cycle_graph():
    # Cycle: A -> B -> C -> A; all nodes have outgoing edges, so none is a 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) # 1.21μs -> 833ns (45.1% faster)

def test_disconnected_graph():
    # Disconnected: A->B, C (isolated); C is a last node (no outgoing edges)
    nodes = [{"id": "A"}, {"id": "B"}, {"id": "C"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 1.50μs -> 708ns (112% faster)

def test_multiple_last_nodes():
    # D->E, F->G, H (isolated); E, G, and H are all last nodes
    nodes = [{"id": "D"}, {"id": "E"}, {"id": "F"}, {"id": "G"}, {"id": "H"}]
    edges = [{"source": "D", "target": "E"}, {"source": "F", "target": "G"}]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 1.54μs -> 750ns (106% faster)

def test_edge_with_nonexistent_node():
    # Edges reference a node not in nodes list; should ignore such edges
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}, {"source": "B", "target": "C"}]  # C does not exist
    # Only B has no outgoing edges among nodes, so B is last
    codeflash_output = find_last_node(nodes, edges) # 1.17μs -> 750ns (55.6% faster)

def test_duplicate_edges():
    # Duplicate edges should not affect result
    nodes = [{"id": "A"}, {"id": "B"}]
    edges = [{"source": "A", "target": "B"}, {"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 708ns (118% faster)

def test_node_with_self_loop():
    # Node with self-loop: A->A; A has outgoing edge, so no last node
    nodes = [{"id": "A"}]
    edges = [{"source": "A", "target": "A"}]
    codeflash_output = find_last_node(nodes, edges) # 1.04μs -> 584ns (78.3% faster)

def test_nodes_with_extra_keys():
    # Nodes have extra keys; function should not care
    nodes = [{"id": "A", "name": "Alpha"}, {"id": "B", "name": "Beta"}]
    edges = [{"source": "A", "target": "B"}]
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 708ns (118% faster)

def test_edge_case_non_string_ids():
    # Node IDs are integers
    nodes = [{"id": 1}, {"id": 2}]
    edges = [{"source": 1, "target": 2}]
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 708ns (118% faster)

def test_edge_case_mixed_type_ids():
    # Node IDs are mixed types; function should handle them as hashable
    nodes = [{"id": 1}, {"id": "2"}]
    edges = [{"source": 1, "target": "2"}]
    codeflash_output = find_last_node(nodes, edges) # 1.54μs -> 709ns (117% faster)

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

def test_large_linear_graph():
    # Linear graph with 1000 nodes: 0->1->2->...->999; last node is 999
    nodes = [{"id": str(i)} for i in range(1000)]
    edges = [{"source": str(i), "target": str(i+1)} for i in range(999)]
    codeflash_output = find_last_node(nodes, edges) # 112μs -> 109μs (2.90% faster)

def test_large_branching_graph():
    # 500 nodes, each node 0..498 has two outgoing edges to i+1 and i+2
    nodes = [{"id": str(i)} for i in range(500)]
    edges = []
    for i in range(498):
        edges.append({"source": str(i), "target": str(i+1)})
        edges.append({"source": str(i), "target": str(i+2)})
    # Last two nodes (498, 499) have no outgoing edges
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 82.6μs -> 79.7μs (3.61% faster)

def test_large_forest_graph():
    # 10 trees of 100 nodes each, each tree is linear
    nodes = [{"id": f"{t}_{i}"} for t in range(10) for i in range(100)]
    edges = []
    for t in range(10):
        for i in range(99):
            edges.append({"source": f"{t}_{i}", "target": f"{t}_{i+1}"})
    # Last node of each tree is a last node
    last_nodes = [{"id": f"{t}_99"} for t in range(10)]
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 67.2μs -> 67.3μs (0.125% slower)

def test_large_graph_with_isolated_nodes():
    # 900 connected nodes in a line, 100 isolated nodes
    nodes = [{"id": str(i)} for i in range(1000)]
    edges = [{"source": str(i), "target": str(i+1)} for i in range(899)]
    # Isolated nodes are 900-999
    last_nodes = [{"id": str(i)} for i in range(900, 1000)]
    last_nodes.append({"id": "899"})  # last node in the line
    codeflash_output = find_last_node(nodes, edges); result = codeflash_output # 104μs -> 100μs (4.33% faster)

def test_large_graph_all_cycles():
    # 100 nodes in a cycle: 0->1->2->...->99->0; no last node
    nodes = [{"id": str(i)} for i in range(100)]
    edges = [{"source": str(i), "target": str((i+1)%100)} for i in range(100)]
    codeflash_output = find_last_node(nodes, edges) # 13.2μs -> 12.5μs (5.34% 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-mc5hogjb and push.

Codeflash

Here is an optimized version of your program.  
### Optimization rationale
- The set comprehension for `source_ids = {e["source"] for e in edges}` is already a good approach.
- However, calling `n["id"] not in source_ids` for every node can be costly, especially if `nodes` is large.  
- The best way to speed this up is to turn `nodes` into an iterable of node IDs for efficient lookup/comparison, but functionally, the code is nearly minimal.
- To slightly speed up.
  - Pre-fetch field lookups to local variables if inside a comprehension.
  - Avoid dict lookup in tight loop by extracting `n["id"]` once.
  - Combine all into a single fast loop without generator expressions.

**That said, the `next((...), None)` daily pattern is already efficient. But you can squeeze out a bit more performance by:**

- Changing the node scan to scan `nodes` only once but precompute the set before.
- Use set subtraction between node IDs and source IDs, and then find *the* node object for the difference.

Below is the code.


**Why is this faster?**

- No per-node `"id"` lookup on every generator iteration.
- Only one pass through nodes, one pass through edges.
- Pure set operations and dict lookups are the fastest way in Python for these container types.

---

If you are guaranteed there is only ever one possible "last" node, you could even do.


But the first (set) version will slightly outperform your original for large lists, as it avoids repeating lookups.

**Choose the version that best suits your typical data size and cardinality.**
@codeflash-ai codeflash-ai bot added the ⚡️ codeflash Optimization PR opened by Codeflash AI label Jun 21, 2025
@codeflash-ai codeflash-ai bot requested a review from KRRT7 June 21, 2025 00:17
@KRRT7 KRRT7 closed this Jun 23, 2025
@codeflash-ai codeflash-ai bot deleted the codeflash/optimize-find_last_node-mc5hogjb branch June 23, 2025 23:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
⚡️ codeflash Optimization PR opened by Codeflash AI
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

1 participant