โ† run

algo-graph-01

1.000
19/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
Looking at this problem, I need to:

1. Find all strongly connected components (SCCs) in the directed graph
2. Contract each SCC into a single node (since once you enter an SCC, you can collect all treasure in it)
3. Build a DAG of the contracted components
4. Find the maximum path sum starting from the component containing `start`

Since gold is non-negative, I want to collect as much as possible, and in a DAG I can use dynamic programming with topological ordering.

Here's my approach:
- Use Kosaraju's algorithm or Tarjan's algorithm to find SCCs (I'll use Kosaraju for clarity)
- Contract the graph: each SCC becomes a node with treasure = sum of treasures in that SCC
- Build the DAG of components
- Perform DP on the DAG in reverse topological order to find maximum path sum from each component
- Return the value for the component containing `start`

```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:
    # Build adjacency list and reverse adjacency list
    graph = [[] for _ in range(n)]
    reverse_graph = [[] for _ in range(n)]
    
    for u, v in roads:
        graph[u].append(v)
        reverse_graph[v].append(u)
    
    # Kosaraju's algorithm to find SCCs
    # Step 1: Do DFS on original graph and get finish order
    visited = [False] * n
    finish_order = []
    
    def dfs_first(u):
        stack = [(u, False)]
        while stack:
            node, processed = stack.pop()
            if processed:
                finish_order.append(node)
                continue
            
            if visited[node]:
                continue
                
            visited[node] = True
            stack.append((node, True))
            
            for v in graph[node]:
                if not visited[v]:
                    stack.append((v, False))
    
    for i in range(n):
        if not visited[i]:
            dfs_first(i)
    
    # Step 2: Do DFS on reverse graph in reverse finish order
    visited = [False] * n
    component_id = [-1] * n
    current_component = 0
    
    def dfs_second(u, comp_id):
        stack = [u]
        while stack:
            node = stack.pop()
            if visited[node]:
                continue
            visited[node] = True
            component_id[node] = comp_id
            
            for v in reverse_graph[node]:
                if not visited[v]:
                    stack.append(v)
    
    # Process in reverse finish order
    for node in reversed(finish_order):
        if not visited[node]:
            dfs_second(node, current_component)
            current_component += 1
    
    num_components = current_component
    
    # Calculate treasure for each component
    component_treasure = [0] * num_components
    for i in range(n):
        comp = component_id[i]
        component_treasure[comp] += treasure[i]
    
    # Build DAG of components
    # Only add edges between different components
    comp_graph = [[] for _ in range(num_components)]
    comp_in_degree = [0] * num_components
    
    for u, v in roads:
        comp_u = component_id[u]
        comp_v = component_id[v]
        if comp_u != comp_v:
            # Avoid duplicate edges
            if comp_v not in comp_graph[comp_u]:
                comp_graph[comp_u].append(comp_v)
                comp_in_degree[comp_v] += 1
    
    # Topological sort using Kahn's algorithm
    queue = deque()
    for i in range(num_components):
        if comp_in_degree[i] == 0:
            queue.append(i)
    
    topo_order = []
    while queue:
        node = queue.popleft()
        topo_order.append(node)
        for neighbor in comp_graph[node]:
            comp_in_degree[neighbor] -= 1
            if comp_in_degree[neighbor] == 0:
                queue.append(neighbor)
    
    # DP: max_treasure_from[comp] = max treasure collectible starting from component comp
    max_treasure_from = [0] * num_components
    
    # Process in reverse topological order (so we process children before parents)
    for comp in reversed(topo_order):
        # Start with the treasure in this component
        max_treasure_from[comp] = component_treasure[comp]
        
        # Try all outgoing edges to see if going to another component gives more
        best_next = 0
        for next_comp in comp_graph[comp]:
            if max_treasure_from[next_comp] > best_next:
                best_next = max_treasure_from[next_comp]
        
        max_treasure_from[comp] += best_next
    
    # Return the max treasure from the component containing start
    start_component = component_id[start]
    return max_treasure_from[start_component]
```

Test output (stdout)

...................                                                      [100%]
19 passed in 0.41s