3331. Find Subtree Sizes After Changes

You are given a tree rooted at node 0 that consists of n nodes numbered from 0 to n - 1. The tree is represented by an array parent of size n, where parent[i] is the parent of node i. Since node 0 is the root, parent[0] == -1.

You are also given a string s of length n, where s[i] is the character assigned to node i.

We make the following changes on the tree one time simultaneously for all nodes x from 1 to n - 1:

• Find the closest node y to node x such that y is an ancestor of x, and s[x] == s[y].
• If node y does not exist, do nothing.
• Otherwise, remove the edge between x and its current parent and make node y the new parent of x by adding an edge between them.

Return an array answer of size n where answer[i] is the size of the subtree rooted at node i in the final tree.

A subtree of treeName is a tree consisting of a node in treeName and all of its descendants.

class Solution {
public:
    vector<int> findSubtreeSizes(vector<int>& parent, string s) {
        vector<int> P = parent;
        for(int i = 1; i < parent.size(); i++) {
            int p = parent[i];
            while(p != -1 and s[p] != s[i]) p = parent[p];
            if(p != -1) P[i] = p;
        }
        vector<int> res(parent.size(), 1), deg(parent.size());
        queue<int> q;
        for(int i = 1; i < parent.size(); i++) deg[P[i]]++;
        for(int i = 1; i < parent.size(); i++) if(!deg[i]) q.push(i);
        while(q.size()) {
            int u = q.front(); q.pop();
            res[P[u]] += res[u];
            if(--deg[P[u]] == 0 and P[u]) q.push(P[u]);
        }
        return res;
    }
};