力扣周賽 5406. 收集樹上所有蘋果的最少時間

給你一棵有 n 個節點的無向樹，節點編號爲 0 到 n-1 ，它們中有一些節點有蘋果。通過樹上的一條邊，需要花費 1 秒鐘。你從 節點 0 出發，請你返回最少需要多少秒，可以收集到所有蘋果，並回到節點 0 。

無向樹的邊由 edges 給出，其中 edges[i] = [fromi, toi] ，表示有一條邊連接 from 和 toi 。除此以外，還有一個布爾數組 hasApple ，其中 hasApple[i] = true 代表節點 i 有一個蘋果，否則，節點 i 沒有蘋果。

 

示例 1：

輸入：n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false,false,true,false,true,true,false]
輸出：8 
解釋：上圖展示了給定的樹，其中紅色節點表示有蘋果。一個能收集到所有蘋果的最優方案由綠色箭頭表示。
示例 2：

輸入：n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false,false,true,false,false,true,false]
輸出：6
解釋：上圖展示了給定的樹，其中紅色節點表示有蘋果。一個能收集到所有蘋果的最優方案由綠色箭頭表示。
示例 3：

輸入：n = 7, edges = [[0,1],[0,2],[1,4],[1,5],[2,3],[2,6]], hasApple = [false,false,false,false,false,false,false]
輸出：0
 

提示：

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

 

class Solution {
public:
    int minTime(int n, vector<vector<int>>& edges, vector<bool>& hasApple) {
        int *fa=new int[n];
        int *ans=new int[n];
        for(int i=0;i<n;i++)ans[i]=0;
        ans[0]=1;
        for(auto it=edges.begin();it!=edges.end();it++)fa[(*it)[1]]=(*it)[0];
        fa[0]=0;
        for(int i=0;i<n;i++)if(hasApple[i])
        {
            int j=i;
            while(ans[j]==0)
            {
                ans[j]=1,j=fa[j];
            }
        }
        int res=-1;
        for(int i=0;i<n;i++)if(ans[i])res++;
        return res*2;
    }
};