Minimum Time to Collect All Apples in a Tree

medium tree dfs hash table depth-first search

Problem

Undirected tree rooted at 0 with n nodes and edges. hasApple[i] tells if node i has an apple. Walking an edge costs 1 second each way. Return the minimum seconds to collect all apples starting and ending at 0.

Inputn=7, edges=[[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple=[false,false,true,false,true,true,false]

Output8

Visit 4 and 5 via 1, and 2 via 0 — total 8 edge traversals (4 round-trips).

edges

hasApple

def min_time(n, edges, hasApple):
    adj = [[] for _ in range(n)]
    for u, v in edges:
        adj[u].append(v); adj[v].append(u)
    def dfs(u, p):
        total = 0
        for v in adj[u]:
            if v == p: continue
            cost = dfs(v, u)
            if cost > 0 or hasApple[v]:
                total += cost + 2
        return total
    return dfs(0, -1)

function minTime(n, edges, hasApple) {
  const adj = Array.from({length: n}, () => []);
  for (const [u, v] of edges) { adj[u].push(v); adj[v].push(u); }
  function dfs(u, p) {
    let total = 0;
    for (const v of adj[u]) {
      if (v === p) continue;
      const cost = dfs(v, u);
      if (cost > 0 || hasApple[v]) total += cost + 2;
    }
    return total;
  }
  return dfs(0, -1);
}

class Solution {
    List> adj;
    boolean[] apple;
    public int minTime(int n, int[][] edges, List hasApple) {
        adj = new ArrayList<>();
        for (int i = 0; i < n; i++) adj.add(new ArrayList<>());
        for (int[] e : edges) { adj.get(e[0]).add(e[1]); adj.get(e[1]).add(e[0]); }
        apple = new boolean[n];
        for (int i = 0; i < n; i++) apple[i] = hasApple.get(i);
        return dfs(0, -1);
    }
    int dfs(int u, int p) {
        int total = 0;
        for (int v : adj.get(u)) {
            if (v == p) continue;
            int cost = dfs(v, u);
            if (cost > 0 || apple[v]) total += cost + 2;
        }
        return total;
    }
}

int minTime(int n, vector>& edges, vector& hasApple) {
    vector> adj(n);
    for (auto& e : edges) { adj[e[0]].push_back(e[1]); adj[e[1]].push_back(e[0]); }
    function dfs = [&](int u, int p) -> int {
        int total = 0;
        for (int v : adj[u]) {
            if (v == p) continue;
            int cost = dfs(v, u);
            if (cost > 0 || hasApple[v]) total += cost + 2;
        }
        return total;
    };
    return dfs(0, -1);
}

Explanation

The key insight is that you only ever need to walk down a branch if there is an apple somewhere inside it. If a whole branch is apple-free, you skip it entirely and save the trip. Each edge you do use costs 2 seconds, because you walk down it and back up.

We turn the edge list into an adjacency list so each node knows its neighbors, then run a DFS from node 0. We pass a parent argument so we never walk back the way we came.

For each child v, we recursively get the cost of its subtree. If that cost is positive (an apple lives deeper) or hasApple[v] is true, then this branch is worth visiting, so we add cost + 2 — the child's own work plus the two seconds for the edge to it.

Example: with apples at nodes 2, 4, 5, the branch to 3 and 6 has no apples and contributes 0. The edges to 4, 5, 1, and 2 each add 2, giving 8 total.

Because each node is visited once and its result bubbles up to its parent, the whole answer is computed in a single DFS sweep.

Time: O(n) Space: O(n)