Minimum Time to Collect All Apples in a Tree

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

class Solution {
public:
    unordered_map<int, vector<int>> g; // to store the graph
    unordered_map<int, bool> visited; // to stop exploring same nodes again and again.
    
    void createGraph(vector<vector<int>>& edges) {
      for (auto e: edges) {
        g[e[0]].push_back(e[1]);
        g[e[1]].push_back(e[0]);
      }
    }
  
    int dfs(int node, int myCost, vector<bool>& hasApple) {
      if (visited[node]) {
          return 0;
      }
      visited[node] = true;
      
      int childrenCost = 0; // cost of traversing all children. 
      for (auto x: g[node]) { 
        childrenCost += dfs(x, 2, hasApple);  // check recursively for all apples.
      }

      if (childrenCost == 0 && hasApple[node] == false) {
        return 0;
      }
      return (childrenCost + myCost);
    }
  
    int minTime(int n, vector<vector<int>>& edges, vector<bool>& hasApple) {
      createGraph(edges); // construct the graph first.
      return dfs(0, 0, hasApple); // cost of reaching the root is 0. For all others, its 2.
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

class Solution {
    public int minTime(int n, int[][] edges, List<Boolean> hasApple) {
        
        Map<Integer, List<Integer>> graph = createGraph(edges); // to store the graph
        Map<Integer, Boolean> visited = new HashMap<>();
		
        return dfs(graph, 0, hasApple, 0, visited); // cost of reaching the root is 0. For all others, its 2.
      
    }
    
    private int dfs(Map<Integer, List<Integer>> graph, int node, List<Boolean> hasApple, int myCost, Map<Integer, Boolean> visited) {
        Boolean v = visited.getOrDefault(node, false);
		if (v) {
			return 0;
		}
		visited.put(node, true);
		
        int childrenCost = 0; // cost of traversing all children. 
      
        for(int n : graph.getOrDefault(node, new ArrayList<>())) {
            childrenCost += dfs(graph, n, hasApple, 2, visited); // check recursively for all apples in subtrees.
        }
      
        if (childrenCost == 0 && hasApple.get(node) == false) {
          // If no child has apples, then we won't traverse the subtree, so cost will be zero.
          // similarly, if current node also does not have the apple, we won't traverse this branch at all, so cost will be zero.
          return 0;
        }
      
        return childrenCost + myCost;
    }
    
    private Map<Integer, List<Integer>> createGraph(int[][] edges) {
        Map<Integer, List<Integer>> graph = new HashMap<>();
      
        for(int i = 0; i < edges.length; i++) {
            List<Integer> list = graph.getOrDefault(edges[i][0], new ArrayList<>()); // Adjecency list representation.
            list.add(edges[i][1]);
            graph.put(edges[i][0], list);
			
			list = graph.getOrDefault(edges[i][1], new ArrayList<>()); // Adjecency list representation.
            list.add(edges[i][0]);
            graph.put(edges[i][1], list);
        }
      
        return graph;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

class Solution:
	def minTime(self, n: int, edges: List[List[int]], hasApple: List[bool]) -> int:
	    adj = [[] for _ in range(n)]
	    for u, v in edges:
	        adj[u].append(v)
	        adj[v].append(u)
	
	    visited = set()
	    def dfs(node):
	        if node in visited:
	            return 0
	        visited.add(node)
	        secs = 0
	        for child in adj[node]:
	            secs += dfs(child)
	        if secs > 0:
	            return secs + 2
	        return 2 if hasApple[node] else 0
	
	    return max(dfs(0) - 2, 0)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

use std::collections::HashMap;

impl Solution {
    pub fn min_time(n: i32, edges: Vec<Vec<i32>>, has_apple: Vec<bool>) -> i32 {
        let mut graph: HashMap<i32, Vec<i32>> = HashMap::new();
        let mut start = 0;
        for i in 0..n {
            graph.insert(i, vec![]);
        }
        for edge in edges.iter() {
            graph.get_mut(&edge[0]).unwrap().push(edge[1]);
            graph.get_mut(&edge[1]).unwrap().push(edge[0]);
        }
        Self::dfs(&graph, &has_apple, 0, -1)
    }

    fn dfs(graph: &HashMap<i32, Vec<i32>>, has_apple: &Vec<bool>, node: i32, prev: i32) -> i32 {
        let mut ans = 0;
        for &nei in graph.get(&node).unwrap() {
            if nei == prev {
                continue;
            }
            let sub_ans = Self::dfs(graph, has_apple, nei, node);
            if sub_ans > 0 || has_apple[nei as usize] {
                ans += sub_ans + 2;
            }
        }
        ans
    }



}