<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>

/*
- classic way: use edges to create {node: [connected nodes]} hash. Then do dfs.
- union-find: 
1. start from node 0: for a connected node, we set roots to be the smallest node among roots.
2. when reach restricted node: just don't join. 
3. at the end, find root and count those with root 0

Wrong 1 (21 mins, 48/62 passed): n =10
edges = [[4,1],[1,3],[1,5],[0,5],[3,6],[8,4],[5,7],[6,9],[3,2]]
restricted = [2,7]
Output 2
Expected 8
Reason: classic mistake: set root of root to new root, instead of setting root to new root!
*/

class Solution {
public:
    template<typename T>
    void printVec(vector<T> cnt) {
        cout << "[ ";
        for(auto& i : cnt) {
            cout << i << " ";
        }
        cout << "]" << endl;
    }

    int find(int node, vector<int>& root) {
        // base case: root is itself
        if(root[node] == node) {
            return node;
        }
        int parent = root[node];
        int root_parent = this->find(parent, root); // find ultimate root of parent
        root[node] = root_parent; // compression of path
        return root_parent;
    }

    int reachableNodes(int n, vector<vector<int>>& edges, vector<int>& restricted) {
        // init root
        vector<int> root(n, 0);
        for(size_t i = 1; i < n; ++i){
            root[i] = i; // init root to itself
        }
        
        unordered_set<int> res_set(restricted.begin(), restricted.end()); // convert to set for fast search
        // iter over edges
        for(auto& edge: edges) {
            int n1 = edge[0];
            int n2 = edge[1];
            if(res_set.count(n1) > 0 || res_set.count(n2) > 0) {
                continue; // do nothing, no union
            }
            // do union
            int root1 = this->find(n1, root);
            int root2 = this->find(n2, root);
            if(root1 == root2) { // already unioned to same root
                continue;
            }
            else if(root1 < root2) {
                root[root2] = root1; // note this is the bug!
            }
            else { // root1 > root2
                root[root1] = root2; // note this is the bug!
            }

            //cout << "after edge (" << n1 << ", " << n2 << "), root vector: ";
            //printVec(root);
        }

        // count those with root 0
        int count = 0;
        for(size_t i = 0; i < n; ++i) {
            int rooti = this->find(i, root);
            if(rooti == 0) {
                count++;
            }
        }
        return count;
    }
};
122940