← run

algo-graph-01

0.474
9/19 tests· algorithms
Challenge · difficulty 4/5
# Treasure runs on a directed road map

An adventurer roams a network of `n` outposts connected by **one-way** roads.
Outpost `i` holds `treasure[i]` gold (a **non-negative** integer). Starting at
outpost `start`, the adventurer may follow roads for as long as they like — even
driving around in circles — and **loots each outpost the first time they visit
it** (revisiting an already-looted outpost yields nothing).

Compute the **maximum total gold** the adventurer can collect.

Implement a file **`solution.py`** with:

```python
def max_treasure(n: int, treasure: Sequence[int],
                 roads: Sequence[Tuple[int, int]], start: int) -> int:
    ...
```

## Input

- `n` — number of outposts, labeled `0 .. n-1`  (`1 <= n`).
- `treasure` — a list of `n` non-negative integers; `treasure[i]` is the gold at outpost `i`.
- `roads` — a list of directed edges `(u, v)` meaning there is a one-way road **from `u` to `v`**.
  Roads may include **self-loops** `(u, u)` and **duplicate** edges.
- `start` — the outpost the adventurer begins at (`0 <= start < n`).

## Output

Return a single integer: the maximum gold collectible on any walk that begins at
`start`, counting each outpost's gold **at most once**.

## What "collect" means

Because roads are one-way, a walk visits a sequence of outposts and loots the
distinct ones it touches. Two consequences drive the problem:

- If a set of outposts is **mutually reachable** (you can get from any one to any
  other and back — a strongly connected group), then once you enter the group you
  can tour **all** of its outposts, so their gold is collected together.
- Once you leave such a group along a road, you can never return to it (otherwise
  it would have been part of the group). So the walk, viewed at the level of these
  groups, moves **strictly forward** through the road network.

Gold is never negative, so lingering to loot more is never harmful; the only real
decision is **which forward branch to commit to**.

## Examples

```python
# 0 -> 1 -> 2, values 1,10,100
assert max_treasure(3, [1, 10, 100], [(0, 1), (1, 2)], 0) == 111
assert max_treasure(3, [1, 10, 100], [(0, 1), (1, 2)], 2) == 100

# 0 <-> 1 is one mutually-reachable group: loot both from either start
assert max_treasure(2, [3, 4], [(0, 1), (1, 0)], 0) == 7

# unreachable fortune is excluded
assert max_treasure(3, [1, 1, 1000], [(0, 1)], 0) == 2

# {0,1,2} form a cycle (1+2+3), then 2 -> 3 worth 100
assert max_treasure(4, [1, 2, 3, 100], [(0, 1), (1, 2), (2, 0), (2, 3)], 0) == 106

# a shared downstream outpost is counted once, not twice
#   0->1->3 and 0->2->3
assert max_treasure(4, [1, 10, 20, 100], [(0, 1), (0, 2), (1, 3), (2, 3)], 0) == 121

# you must pick the richer of two mutually-exclusive branches
assert max_treasure(3, [1, 5, 50], [(0, 1), (0, 2)], 0) == 51
```

## Constraints & notes

- Only outposts **reachable from `start`** can ever be looted.
- Self-loops and duplicate roads must be handled gracefully.
- The graph can be **large and deep**: expect up to about `10^5` outposts and
  `2 * 10^5` roads, including a single chain that long. A solution whose recursion
  depth grows with the graph will overflow — use an **iterative** approach. An
  overall `O(n + m)` algorithm is expected.
tests/test_treasure.py
import random

from solution import max_treasure


# --------------------------------------------------------------------------
# Independent reference: O(n^3) transitive-closure SCCs + recursive DAG DP.
# Deliberately a different algorithm than the solution (no Kosaraju/Tarjan).
# Only used on small graphs.
# --------------------------------------------------------------------------
def _brute(n, treasure, roads, start):
    adj = [[] for _ in range(n)]
    for u, v in roads:
        adj[u].append(v)

    reach = [[False] * n for _ in range(n)]
    for s in range(n):
        seen = {s}
        stack = [s]
        while stack:
            x = stack.pop()
            for y in adj[x]:
                if y not in seen:
                    seen.add(y)
                    stack.append(y)
        for y in seen:
            reach[s][y] = True
        reach[s][s] = True

    comp = [-1] * n
    cid = 0
    for i in range(n):
        if comp[i] != -1:
            continue
        comp[i] = cid
        for j in range(n):
            if comp[j] == -1 and reach[i][j] and reach[j][i]:
                comp[j] = cid
        cid += 1

    cw = [0] * cid
    for i in range(n):
        cw[comp[i]] += treasure[i]

    cadj = [set() for _ in range(cid)]
    for u, v in roads:
        if comp[u] != comp[v]:
            cadj[comp[u]].add(comp[v])

    memo = {}

    def dp(cc):
        if cc in memo:
            return memo[cc]
        best = 0
        for sc in cadj[cc]:
            best = max(best, dp(sc))
        memo[cc] = cw[cc] + best
        return memo[cc]

    return dp(comp[start])


# --------------------------- basic / edge cases ---------------------------
def test_single_node_no_edges():
    assert max_treasure(1, [7], [], 0) == 7


def test_single_node_with_self_loop():
    assert max_treasure(1, [7], [(0, 0)], 0) == 7


def test_zero_value_node():
    assert max_treasure(1, [0], [], 0) == 0


def test_simple_chain():
    # 0 -> 1 -> 2, collect everything downstream
    treasure = [1, 10, 100]
    roads = [(0, 1), (1, 2)]
    assert max_treasure(3, treasure, roads, 0) == 111
    assert max_treasure(3, treasure, roads, 1) == 110
    assert max_treasure(3, treasure, roads, 2) == 100


def test_two_node_cycle_is_one_scc():
    # 0 <-> 1: both lootable from either start
    assert max_treasure(2, [3, 4], [(0, 1), (1, 0)], 0) == 7
    assert max_treasure(2, [3, 4], [(0, 1), (1, 0)], 1) == 7


def test_unreachable_treasure_excluded():
    # node 2 holds a fortune but is unreachable from 0
    treasure = [1, 1, 1000]
    roads = [(0, 1)]
    assert max_treasure(3, treasure, roads, 0) == 2


def test_pick_the_richer_branch():
    # 0 -> 1 (value 5) and 0 -> 2 (value 50); only one branch is on a path
    treasure = [1, 5, 50]
    roads = [(0, 1), (0, 2)]
    assert max_treasure(3, treasure, roads, 0) == 51


def test_branch_choice_favors_longer_sum():
    # 0 -> 1 -> 3 (1+2+100) vs 0 -> 2 (1+50) ; note 3 unreachable via 2
    treasure = [1, 2, 50, 100]
    roads = [(0, 1), (1, 3), (0, 2)]
    assert max_treasure(4, treasure, roads, 0) == 103


def test_scc_then_dag_tail():
    # {0,1,2} form a cycle (sum 6), then 2 -> 3 (value 100)
    treasure = [1, 2, 3, 100]
    roads = [(0, 1), (1, 2), (2, 0), (2, 3)]
    assert max_treasure(4, treasure, roads, 0) == 106
    assert max_treasure(4, treasure, roads, 3) == 100


def test_multi_edges_and_self_loops_ignored():
    treasure = [10, 20]
    roads = [(0, 1), (0, 1), (0, 0), (1, 1)]
    assert max_treasure(2, treasure, roads, 0) == 30


def test_diamond_shared_tail_counted_once():
    #   0 -> 1 -> 3
    #   0 -> 2 -> 3
    # node 3 must be counted once, not twice
    treasure = [1, 10, 20, 100]
    roads = [(0, 1), (0, 2), (1, 3), (2, 3)]
    assert max_treasure(4, treasure, roads, 0) == 1 + 20 + 100  # via 0->2->3


def test_start_in_middle_ignores_upstream():
    treasure = [1000, 1, 1]
    roads = [(0, 1), (1, 2)]
    # starting at 1, the rich node 0 is upstream and unreachable
    assert max_treasure(3, treasure, roads, 1) == 2


def test_two_cycles_joined():
    # cycle A {0,1} -> cycle B {2,3}
    treasure = [1, 1, 5, 5]
    roads = [(0, 1), (1, 0), (1, 2), (2, 3), (3, 2)]
    assert max_treasure(4, treasure, roads, 0) == 12
    assert max_treasure(4, treasure, roads, 2) == 10


# --------------------------- randomized fuzz ---------------------------
def test_random_small_graphs_match_brute():
    rng = random.Random(20260701)
    for _ in range(600):
        n = rng.randint(1, 8)
        treasure = [rng.randint(0, 12) for _ in range(n)]
        m = rng.randint(0, n * 2)
        roads = [(rng.randrange(n), rng.randrange(n)) for _ in range(m)]
        start = rng.randrange(n)
        expected = _brute(n, treasure, roads, start)
        assert max_treasure(n, treasure, roads, start) == expected, (
            n,
            treasure,
            roads,
            start,
        )


def test_random_dense_graphs_match_brute():
    rng = random.Random(4242)
    for _ in range(300):
        n = rng.randint(3, 7)
        treasure = [rng.randint(0, 20) for _ in range(n)]
        roads = []
        for u in range(n):
            for v in range(n):
                if u != v and rng.random() < 0.4:
                    roads.append((u, v))
        start = rng.randrange(n)
        expected = _brute(n, treasure, roads, start)
        assert max_treasure(n, treasure, roads, start) == expected


# --------------------------- larger / performance ---------------------------
def test_large_chain():
    n = 100_000
    treasure = [1] * n
    roads = [(i, i + 1) for i in range(n - 1)]
    assert max_treasure(n, treasure, roads, 0) == n
    assert max_treasure(n, treasure, roads, n // 2) == n - n // 2


def test_large_single_cycle():
    n = 100_000
    treasure = [1] * n
    roads = [(i, (i + 1) % n) for i in range(n)]
    # entire cycle is one SCC, all lootable from anywhere
    assert max_treasure(n, treasure, roads, 0) == n
    assert max_treasure(n, treasure, roads, 12345) == n


def test_large_deep_no_recursion_limit():
    # A long chain would overflow a naive recursive SCC / DP; must be iterative.
    n = 50_000
    treasure = [2] * n
    roads = [(i, i + 1) for i in range(n - 1)]
    assert max_treasure(n, treasure, roads, 0) == 2 * n


def test_large_binary_tree_longest_root_to_leaf():
    # complete binary tree, node i -> 2i+1, 2i+2; value = 1 everywhere.
    # longest root path collects (deepest level + 1) nodes.
    n = 1 << 16  # 65536 nodes
    treasure = [1] * n
    roads = []
    for i in range(n):
        for ch in (2 * i + 1, 2 * i + 2):
            if ch < n:
                roads.append((i, ch))
    # expected = number of levels on the deepest greedy root-to-leaf walk
    levels = 0
    idx = 0
    while idx < n:
        levels += 1
        idx = 2 * idx + 1
    assert max_treasure(n, treasure, roads, 0) == levels
System prompt
You are an expert programmer. Solve the task exactly as specified. Output your solution as fenced code blocks using the required file name(s) and the exact function/type signatures requested. Prefer correctness; do not include prose outside code unless asked.
Proposed solution
```python
from typing import Sequence, Tuple
from collections import defaultdict, deque

def max_treasure(n: int, treasure: Sequence[int], roads: Sequence[Tuple[int, int]], start: int) -> int:
    graph = defaultdict(list)
    for u, v in roads:
        graph[u].append(v)

    visited = [False] * n
    max_gold = 0

    def dfs(outpost):
        nonlocal max_gold
        visited[outpost] = True
        total_gold = treasure[outpost]
        for neighbor in graph[outpost]:
            if not visited[neighbor]:
                total_gold += dfs(neighbor)
        max_gold = max(max_gold, total_gold)
        visited[outpost] = False
        return total_gold

    total_gold = dfs(start)
    return total_gold
```

Test output (stdout)

^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
solution.py:18: in dfs
    total_gold += dfs(neighbor)
                  ^^^^^^^^^^^^^
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

outpost = 965

    def dfs(outpost):
        nonlocal max_gold
        visited[outpost] = True
        total_gold = treasure[outpost]
        for neighbor in graph[outpost]:
            if not visited[neighbor]:
>               total_gold += dfs(neighbor)
                              ^^^^^^^^^^^^^
E               RecursionError: maximum recursion depth exceeded

solution.py:18: RecursionError
_________________ test_large_binary_tree_longest_root_to_leaf __________________

    def test_large_binary_tree_longest_root_to_leaf():
        # complete binary tree, node i -> 2i+1, 2i+2; value = 1 everywhere.
        # longest root path collects (deepest level + 1) nodes.
        n = 1 << 16  # 65536 nodes
        treasure = [1] * n
        roads = []
        for i in range(n):
            for ch in (2 * i + 1, 2 * i + 2):
                if ch < n:
                    roads.append((i, ch))
        # expected = number of levels on the deepest greedy root-to-leaf walk
        levels = 0
        idx = 0
        while idx < n:
            levels += 1
            idx = 2 * idx + 1
>       assert max_treasure(n, treasure, roads, 0) == levels
E       assert 65536 == 17
E        +  where 65536 = max_treasure(65536, [1, 1, 1, 1, 1, 1, ...], [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6), ...], 0)

test_treasure.py:226: AssertionError
=========================== short test summary info ============================
FAILED test_treasure.py::test_pick_the_richer_branch - assert 56 == 51
FAILED test_treasure.py::test_branch_choice_favors_longer_sum - assert 153 ==...
FAILED test_treasure.py::test_multi_edges_and_self_loops_ignored - assert 50 ...
FAILED test_treasure.py::test_diamond_shared_tail_counted_once - assert 231 =...
FAILED test_treasure.py::test_random_small_graphs_match_brute - AssertionErro...
FAILED test_treasure.py::test_random_dense_graphs_match_brute - assert 116 == 57
FAILED test_treasure.py::test_large_chain - RecursionError: maximum recursion...
FAILED test_treasure.py::test_large_single_cycle - RecursionError: maximum re...
FAILED test_treasure.py::test_large_deep_no_recursion_limit - RecursionError:...
FAILED test_treasure.py::test_large_binary_tree_longest_root_to_leaf - assert...
10 failed, 9 passed in 0.51s
Self-repair attempts
[
  {
    "total": 19,
    "answer": "```python\nfrom typing import Sequence, Tuple\nfrom collections import defaultdict, deque\n\ndef max_treasure(n: int, treasure: Sequence[int], roads: Sequence[Tuple[int, int]], start: int) -> int:\n    graph = defaultdict(list)\n    for u, v in roads:\n        graph[u].append(v)\n\n    visited = [False] * n\n    max_gold = 0\n\n    def dfs(outpost):\n        nonlocal max_gold\n        visited[outpost] = True\n        total_gold = treasure[outpost]\n        for neighbor in graph[outpost]:\n            if not visited[neighbor]:\n                total_gold += dfs(neighbor)\n        max_gold = max(max_gold, total_gold)\n        visited[outpost] = False\n        return total_gold\n\n    total_gold = dfs(start)\n    return total_gold\n```",
    "passed": 9,
    "test_error": "n.py:18: in dfs\n    total_gold += dfs(neighbor)\n                  ^^^^^^^^^^^^^\nsolution.py:18: in dfs\n    total_gold += dfs(neighbor)\n                  ^^^^^^^^^^^^^\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \n\noutpost = 965\n\n    def dfs(outpost):\n        nonlocal max_gold\n        visited[outpost] = True\n        total_gold = treasure[outpost]\n        for neighbor in graph[outpost]:\n            if not visited[neighbor]:\n>               total_gold += dfs(neighbor)\n                              ^^^^^^^^^^^^^\nE               RecursionError: maximum recursion depth exceeded\n\nsolution.py:18: RecursionError\n_________________ test_large_binary_tree_longest_root_to_leaf __________________\n\n    def test_large_binary_tree_longest_root_to_leaf():\n        # complete binary tree, node i -> 2i+1, 2i+2; value = 1 everywhere.\n        # longest root path collects (deepest level + 1) nodes.\n        n = 1 << 16  # 65536 nodes\n        treasure = [1] * n\n        roads = []\n        for i in range(n):\n            for ch in (2 * i + 1, 2 * i + 2):\n                if ch < n:\n                    roads.append((i, ch))\n        # expected = number of levels on the deepest greedy root-to-leaf walk\n        levels = 0\n        idx = 0\n        while idx < n:\n            levels += 1\n            idx = 2 * idx + 1\n>       assert max_treasure(n, treasure, roads, 0) == levels\nE       assert 65536 == 17\nE        +  where 65536 = max_treasure(65536, [1, 1, 1, 1, 1, 1, ...], [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6), ...], 0)\n\ntest_treasure.py:226: AssertionError\n=========================== short test summary info ============================\nFAILED test_treasure.py::test_pick_the_richer_branch - assert 56 == 51\nFAILED test_treasure.py::test_branch_choice_favors_longer_sum - assert 153 ==...\nFAILED test_treasure.py::test_multi_edges_and_self_loops_ignored - assert 50 ...\nFAILED test_treasure.py::test_diamond_shared_tail_counted_once - assert 231 =...\nFAILED test_treasure.py::test_random_small_graphs_match_brute - AssertionErro...\nFAILED test_treasure.py::test_random_dense_graphs_match_brute - assert 116 == 57\nFAILED test_treasure.py::test_large_chain - RecursionError: maximum recursion...\nFAILED test_treasure.py::test_large_single_cycle - RecursionError: maximum re...\nFAILED test_treasure.py::test_large_deep_no_recursion_limit - RecursionError:...\nFAILED test_treasure.py::test_large_binary_tree_longest_root_to_leaf - assert...\n10 failed, 9 passed in 0.51s"
  },
  {
    "total": 19,
    "answer": "```python\nfrom typing import Sequence, Tuple, List\nfrom collections import defaultdict, deque\n\ndef max_treasure(n: int, treasure: Sequence[int], roads: Sequence[Tuple[int, int]], start: int) -> int:\n    graph = defaultdict(list)\n    for u, v in roads:\n        graph[u].append(v)\n\n    visited = [False] * n\n    max_gold = [0]\n\n    def dfs(outpost):\n        visited[outpost] = True\n        total_gold = treasure[outpost]\n        for neighbor in graph[outpost]:\n            if not visited[neighbor]:\n                total_gold += dfs(neighbor)\n        max_gold[0] = max(max_gold[0], total_gold)\n        visited[outpost] = False\n        return total_gold\n\n    dfs(start)\n    return max_gold[0]\n```",
    "passed": 9,
    "test_error": "     ^^^^^^^^^^^^^\nsolution.py:17: in dfs\n    total_gold += dfs(neighbor)\n                  ^^^^^^^^^^^^^\nsolution.py:17: in dfs\n    total_gold += dfs(neighbor)\n                  ^^^^^^^^^^^^^\n_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ \n\noutpost = 965\n\n    def dfs(outpost):\n        visited[outpost] = True\n        total_gold = treasure[outpost]\n        for neighbor in graph[outpost]:\n            if not visited[neighbor]:\n>               total_gold += dfs(neighbor)\n                              ^^^^^^^^^^^^^\nE               RecursionError: maximum recursion depth exceeded\n\nsolution.py:17: RecursionError\n_________________ test_large_binary_tree_longest_root_to_leaf __________________\n\n    def test_large_binary_tree_longest_root_to_leaf():\n        # complete binary tree, node i -> 2i+1, 2i+2; value = 1 everywhere.\n        # longest root path collects (deepest level + 1) nodes.\n        n = 1 << 16  # 65536 nodes\n        treasure = [1] * n\n        roads = []\n        for i in range(n):\n            for ch in (2 * i + 1, 2 * i + 2):\n                if ch < n:\n                    roads.append((i, ch))\n        # expected = number of levels on the deepest greedy root-to-leaf walk\n        levels = 0\n        idx = 0\n        while idx < n:\n            levels += 1\n            idx = 2 * idx + 1\n>       assert max_treasure(n, treasure, roads, 0) == levels\nE       assert 65536 == 17\nE        +  where 65536 = max_treasure(65536, [1, 1, 1, 1, 1, 1, ...], [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6), ...], 0)\n\ntest_treasure.py:226: AssertionError\n=========================== short test summary info ============================\nFAILED test_treasure.py::test_pick_the_richer_branch - assert 56 == 51\nFAILED test_treasure.py::test_branch_choice_favors_longer_sum - assert 153 ==...\nFAILED test_treasure.py::test_multi_edges_and_self_loops_ignored - assert 50 ...\nFAILED test_treasure.py::test_diamond_shared_tail_counted_once - assert 231 =...\nFAILED test_treasure.py::test_random_small_graphs_match_brute - AssertionErro...\nFAILED test_treasure.py::test_random_dense_graphs_match_brute - assert 116 == 57\nFAILED test_treasure.py::test_large_chain - RecursionError: maximum recursion...\nFAILED test_treasure.py::test_large_single_cycle - RecursionError: maximum re...\nFAILED test_treasure.py::test_large_deep_no_recursion_limit - RecursionError:...\nFAILED test_treasure.py::test_large_binary_tree_longest_root_to_leaf - assert...\n10 failed, 9 passed in 0.52s"
  },
  {
    "total": 19,
    "answer": "```python\nfrom typing import Sequence, Tuple, List\nfrom collections import defaultdict, deque\n\ndef max_treasure(n: int, treasure: Sequence[int], roads: Sequence[Tuple[int, int]], start: int) -> int:\n    graph = defaultdict(list)\n    for u, v in roads:\n        graph[u].append(v)\n\n    visited = [False] * n\n    max_gold = [0]\n\n    def dfs(outpost):\n        visited[outpost] = True\n        total_gold = treasure[outpost]\n        for neighbor in graph[outpost]:\n            if not visited[neighbor]:\n                total_gold += dfs(neighbor)\n        max_gold[0] = max(max_gold[0], total_gold)\n        visited[outpost] = False\n        return total_gold\n\n    dfs(start)\n    return max_gold[0]\n```",
    "passed": 9,
    "test_error": ""
  }
]