<p>There is an undirected tree with <code>n</code> nodes labeled from <code>0</code> to <code>n - 1</code> and <code>n - 1</code> edges.</p>

<p>You are given a 2D integer array <code>edges</code> of length <code>n - 1</code> where <code>edges[i] = [a<sub>i</sub>, b<sub>i</sub>]</code> indicates that there is an edge between nodes <code>a<sub>i</sub></code> and <code>b<sub>i</sub></code> in the tree. You are also given an integer array <code>restricted</code> which represents <strong>restricted</strong> nodes.</p>

<p>Return <em>the <strong>maximum</strong> number of nodes you can reach from node </em><code>0</code><em> without visiting a restricted node.</em></p>

<p>Note that node <code>0</code> will <strong>not</strong> be a restricted node.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2022/06/15/ex1drawio.png" style="width: 402px; height: 322px;" />
<pre>
<strong>Input:</strong> n = 7, edges = [[0,1],[1,2],[3,1],[4,0],[0,5],[5,6]], restricted = [4,5]
<strong>Output:</strong> 4
<strong>Explanation:</strong> The diagram above shows the tree.
We have that [0,1,2,3] are the only nodes that can be reached from node 0 without visiting a restricted node.
</pre>

<p><strong class="example">Example 2:</strong></p>
<img alt="" src="https://assets.leetcode.com/uploads/2022/06/15/ex2drawio.png" style="width: 412px; height: 312px;" />
<pre>
<strong>Input:</strong> n = 7, edges = [[0,1],[0,2],[0,5],[0,4],[3,2],[6,5]], restricted = [4,2,1]
<strong>Output:</strong> 3
<strong>Explanation:</strong> The diagram above shows the tree.
We have that [0,5,6] are the only nodes that can be reached from node 0 without visiting a restricted node.
</pre>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>2 &lt;= n &lt;= 10<sup>5</sup></code></li>
	<li><code>edges.length == n - 1</code></li>
	<li><code>edges[i].length == 2</code></li>
	<li><code>0 &lt;= a<sub>i</sub>, b<sub>i</sub> &lt; n</code></li>
	<li><code>a<sub>i</sub> != b<sub>i</sub></code></li>
	<li><code>edges</code> represents a valid tree.</li>
	<li><code>1 &lt;= restricted.length &lt; n</code></li>
	<li><code>1 &lt;= restricted[i] &lt; n</code></li>
	<li>All the values of <code>restricted</code> are <strong>unique</strong>.</li>
</ul>

/*
Important to note that when doing Union, it is not necessary that it gets unioned to node 0 only. So when want size of node 0, get parent of 0 and then use it's size.

DSU was my initial approach

When using DFS, no need for extra restricted set/array. Just mark it as visited.
*/
/*
class Solution {
    int dfs(vector<vector<int>>& adj, vector<bool>& visitedOrRestricted, int node){
        visitedOrRestricted[node]=true;
        int curr=1;
        for(int child:adj[node]){
            if(!visitedOrRestricted[child])
                curr+=dfs(adj,visitedOrRestricted,child);
        }
        return curr;
    }
public:
    int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
        vector<bool> visitedOrRestricted(n,false);
        for(int &node:restricted)
            visitedOrRestricted[node]=true;
        vector<vector<int>> adj(n);
        for(auto &edge:edges){
            int u=edge[0], v=edge[1];
            adj[u].push_back(v);
            adj[v].push_back(u);
        }
        return dfs(adj,visitedOrRestricted,0);
    }
};
*/
// DSU
class DSU {
    public:
    vector<int> parent,size;
    DSU(int n){
        parent.resize(n);
        iota(parent.begin(),parent.end(),0);
        size.resize(n,1);
    }
    int Find(int v){
        if(parent[v]==v) return v;
        return parent[v] = Find(parent[v]);
    }
    void Union(int a, int b){
        a=Find(a);
        b=Find(b);
        if(a==b) return;
        if(size[a]<size[b]) swap(a,b);
        parent[b]=a;
        size[a]+=size[b];
    }
};
class Solution {
public:
    int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
        DSU dsu(n);
        // unordered_set<int> hset;
        // for(int &node:restricted)
            // hset.insert(node);
        unordered_set<int> hset(restricted.begin(), restricted.end());
        for(auto &edge:edges){
            int u=edge[0], v=edge[1];
            // if(hset.count(u) || hset.count(v)) continue;
            if(hset.count(u) + hset.count(v) == 0)
                dsu.Union(u,v);
        }
        // int parentOfZero = dsu.parent[0];
        // return dsu.size[parentOfZero==-1 ? 0 : parentOfZero];
        return dsu.size[ dsu.parent[0] ];
    }
};

126900