LeetCode 1443 Minimum Time to Collect All Apples in a Tree (dfs)

Given an undirected tree consisting of n vertices numbered from 0 to n-1, which has some apples in their vertices. You spend 1 second to walk over one edge of the tree. Return the minimum time in seconds you have to spend to collect all apples in the tree, starting at vertex 0 and coming back to this vertex.

The edges of the undirected tree are given in the array edges, where edges[i] = [ai, bi] means that exists an edge connecting the vertices ai and bi. Additionally, there is a boolean array hasApple, where hasApple[i] = true means that vertex i has an apple; otherwise, it does not have any apple.

 

Example 1:

Input: n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false,false,true,false,true,true,false]
Output: 8 
Explanation: The figure above represents the given tree where red vertices have an apple. One optimal path to collect all apples is shown by the green arrows.  

Example 2:

Input: n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false,false,true,false,false,true,false]
Output: 6
Explanation: The figure above represents the given tree where red vertices have an apple. One optimal path to collect all apples is shown by the green arrows.  

Example 3:

Input: n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false,false,false,false,false,false,false]
Output: 0

 

Constraints:

• 1 <= n <= 10^5
• edges.length == n - 1
• edges[i].length == 2
• 0 <= ai < bi <= n - 1
• fromi < toi
• hasApple.length == n

题目链接：https://leetcode.com/problems/minimum-time-to-collect-all-apples-in-a-tree/

题目大意：红点是苹果，求从根（0）出发采完所有苹果并回到根所走路程的最小值

题目分析：很明显的分治，一开始的思路是判断以每个节点为根的子树中是否含有苹果，若含有则从根到其的这段路是必经的，故先预处理一下然后dfs累加即可，最后减2目的是消除根节点的影响

时间17ms，击败98.29%

内存73.6MB，击败98.57%

class Solution {
    
    class Edge {
        int to, nxt;
        Edge(int to, int nxt) {
            this.to = to;
            this.nxt = nxt;
        }
    }
    
    class AppleTree {
        int n, m, cnt;
        int[] head;
        Edge[] e;
        List<Boolean> hasApple;
        boolean[] subTreeHasApple;
        boolean[] vis;
        AppleTree(int n, int m, List<Boolean> hasApple) {
            this.n = n;
            this.m = m;
            this.e = new Edge[m << 1];
            this.head = new int[n];
            Arrays.fill(this.head, -1);
            this.cnt = 0;
            this.hasApple = hasApple;
            this.subTreeHasApple = new boolean[n];
            this.vis = new boolean[n];
        }

        void clearVis() {
            Arrays.fill(this.vis, false);
        }
        
        void addEdge(int u, int v) {
            e[cnt] = new Edge(v, head[u]);
            head[u] = cnt++;
        }
        
        boolean findApple(int u) {
            vis[u] = true;
            boolean ok = this.hasApple.get(u);
            for (int i = head[u]; i != -1; i = e[i].nxt) {
                int v = e[i].to;
                if (v != u && !vis[v]) {
                    ok |= findApple(v);
                }
            }
            this.subTreeHasApple[u] = ok;
            return ok;
        }
        
        int solve(int u) {
            vis[u] = true;
            if (this.subTreeHasApple[u]) {
                int cur = 2;
                for (int i = head[u]; i != -1; i = e[i].nxt) {
                    int v = e[i].to;
                    if (u != v && !vis[v]) {
                        cur += solve(v);
                    }
                }
                return cur;   
            }
            return 0;
        }
    }
    
    public int minTime(int n, int[][] edges, List<Boolean> hasApple) {
        if (n == 1) {
            return 0;
        }
        AppleTree at = new AppleTree(n, edges.length, hasApple);
        for (int i = 0; i < edges.length; i++) {
            at.addEdge(edges[i][0], edges[i][1]);
            at.addEdge(edges[i][1], edges[i][0]);
        }
        at.findApple(0);
        // for (int i = 0; i < n; i++) {
        //     System.out.println("node[" + i + "] = " + at.subTreeHasApple[i]);
        // }
        at.clearVis();
        int res = at.solve(0);
        if (res > 0) {
            res -= 2;
        }
        return res;
    }
}

进一步发现预处理和计算的过程可以合并

时间12ms，击败98.29%

内存73.2MB，击败99.43%

class Solution {
    
    class Edge {
        int to, nxt;
        Edge(int to, int nxt) {
            this.to = to;
            this.nxt = nxt;
        }
    }
    
    class AppleTree {
        int n, m, cnt;
        int[] head;
        Edge[] e;
        List<Boolean> hasApple;
        boolean[] vis;
        int ans;
        AppleTree(int n, int m, List<Boolean> hasApple) {
            this.n = n;
            this.m = m;
            this.e = new Edge[m << 1];
            this.head = new int[n];
            Arrays.fill(this.head, -1);
            this.cnt = 0;
            this.hasApple = hasApple;
            this.vis = new boolean[n];
            this.ans = 0;
        }
        
        void addEdge(int u, int v) {
            e[cnt] = new Edge(v, head[u]);
            head[u] = cnt++;
        }
        
        boolean solve(int u) {
            vis[u] = true;
            boolean ok = this.hasApple.get(u);
            for (int i = head[u]; i != -1; i = e[i].nxt) {
                int v = e[i].to;
                if (!vis[v]) {
                    ok |= solve(v);
                }
            }
            if (ok) {
                ans += 2;
            }
            return ok;
        }
    }
    
    public int minTime(int n, int[][] edges, List<Boolean> hasApple) {
        if (n == 1) {
            return 0;
        }
        AppleTree at = new AppleTree(n, edges.length, hasApple);
        for (int i = 0; i < edges.length; i++) {
            at.addEdge(edges[i][0], edges[i][1]);
            at.addEdge(edges[i][1], edges[i][0]);
        }
        at.solve(0);
        if (at.ans > 0) {
            at.ans -= 2;
        }
        return at.ans;
    }
}