A tree problem: Special Nodes

This is a pretty hard problem.

Special Nodes

Naively solve the problem by using a dfs way to traverse each subtree.

Each subtree will maintain a set contains all the colors values. Then check each subtree by using a DFS way.

int ret = 0;
int seen[100010];
bool helper(int node, vector<vector<int>>& tree, vector<int>& color, unordered_set<int>&s) {
    unordered_set<int> thiss;
    auto x = color[node];
    s.insert(x);
    bool good = true;
    for (auto &next: tree[node]) {
        if (!seen[next]) {
            thiss.clear();
            seen[next] = 1;
            auto good_ = helper(next, tree, color, thiss);
            if (!good_) {
                s.clear();
                good = false;
            } else {
                // magic line to get rid of TLE
                if (thiss.size()> s.size()) swap(s, thiss);
                for (auto item: thiss){
                    if (s.find(item) != s.end()){good = false;}
                    else s.insert(item);
                }
            }
        }
    }
    ret += good;
    return good;
}
int solve(vector<vector<int>>& tree, vector<int>& color) {
    ret = 0;
    memset(seen, 0, sizeof seen);
    unordered_set<int> s;
    seen[0] = 1;
    auto good = helper(0, tree, color, s);
    return ret; 
}

We can have a better O(n) solution to solve this problem. The solution is pretty smart. See the post here.