algo-graph-01
1.000
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
To solve this problem, we need to find the maximum total gold an adventurer can collect starting from a given outpost in a directed graph. The key observation is that the adventurer can collect all the gold in any strongly connected component (SCC) they enter, and once they leave an SCC, they cannot return. Therefore, the problem reduces to finding the maximum sum of gold in the SCCs reachable from the start node, considering the dependencies between these SCCs.
### Approach
1. **Identify Strongly Connected Components (SCCs)**: Using Kosaraju's algorithm or Tarjan's algorithm, we can decompose the graph into SCCs. Each SCC is a maximal subgraph where every node is reachable from every other node in the subgraph.
2. **Condense the Graph**: Treat each SCC as a single node in a new condensed graph. The edges between these nodes represent the original edges between different SCCs.
3. **Topological Sorting**: Perform a topological sort on the condensed graph to process the SCCs in an order where all dependencies (incoming edges) are resolved before processing a node.
4. **Dynamic Programming (DP) on Condensed Graph**: For each SCC in the topological order, compute the maximum gold that can be collected starting from that SCC. The DP value for an SCC is the sum of its own gold plus the maximum DP value of any SCC reachable from it (since once you leave the current SCC, you can choose the best path forward).
### Solution Code
```python
from typing import Sequence, Tuple
from collections import deque
def max_treasure(n: int, treasure: Sequence[int], roads: Sequence[Tuple[int, int]], start: int) -> int:
# Step 1: Build the graph and reverse graph for Kosaraju's algorithm
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)
# Step 2: First pass of Kosaraju's algorithm to get finishing times
visited = [False] * n
order = []
def dfs(u):
stack = [(u, False)]
while stack:
node, processed = stack.pop()
if processed:
order.append(node)
continue
if visited[node]:
continue
visited[node] = True
stack.append((node, True))
for neighbor in graph[node]:
if not visited[neighbor]:
stack.append((neighbor, False))
for i in range(n):
if not visited[i]:
dfs(i)
# Step 3: Second pass on the reverse graph to find SCCs
visited = [False] * n
scc_id = [0] * n
current_id = 0
scc_treasure = [0] * n # This will store the sum of treasures in each SCC
def reverse_dfs(u):
stack = [u]
visited[u] = True
scc_treasure[current_id] += treasure[u]
scc_id[u] = current_id
while stack:
node = stack.pop()
for neighbor in reverse_graph[node]:
if not visited[neighbor]:
visited[neighbor] = True
scc_treasure[current_id] += treasure[neighbor]
scc_id[neighbor] = current_id
stack.append(neighbor)
for node in reversed(order):
if not visited[node]:
reverse_dfs(node)
current_id += 1
# Step 4: Build the condensed graph
condensed_graph = [[] for _ in range(current_id)]
in_degree = [0] * current_id
for u, v in roads:
u_scc = scc_id[u]
v_scc = scc_id[v]
if u_scc != v_scc:
condensed_graph[u_scc].append(v_scc)
in_degree[v_scc] += 1
# Step 5: Topological sort on the condensed graph
queue = deque()
for i in range(current_id):
if in_degree[i] == 0:
queue.append(i)
topo_order = []
while queue:
u = queue.popleft()
topo_order.append(u)
for v in condensed_graph[u]:
in_degree[v] -= 1
if in_degree[v] == 0:
queue.append(v)
# Step 6: Dynamic programming to compute max treasure for each SCC
dp = [0] * current_id
for u in reversed(topo_order):
max_next = 0
for v in condensed_graph[u]:
if dp[v] > max_next:
max_next = dp[v]
dp[u] = scc_treasure[u] + max_next
# The start node's SCC is scc_id[start]
return dp[scc_id[start]]
```
### Explanation
1. **Graph Construction**: The graph and its reverse are built from the given roads to facilitate Kosaraju's algorithm.
2. **Kosaraju's Algorithm**: This algorithm is used to find SCCs by performing a DFS on the original graph to determine the order of nodes and then performing a DFS on the reverse graph in reverse order to identify SCCs.
3. **Condensed Graph**: The original graph is condensed into a graph where each node represents an SCC. Edges between these nodes represent the original edges between different SCCs.
4. **Topological Sorting**: The condensed graph is topologically sorted to process nodes in an order where all dependencies are resolved before processing a node.
5. **Dynamic Programming**: For each SCC in reverse topological order, the maximum gold that can be collected is computed as the sum of the gold in the SCC plus the maximum gold from any reachable SCC. This ensures that once an SCC is processed, all possible paths forward have been considered.
This approach efficiently handles the problem by leveraging SCC decomposition and dynamic programming on the condensed graph, ensuring optimal performance even for large graphs.Test output (stdout)
................... [100%] 19 passed in 0.40s