Tree problem

1519. Number of Nodes in the Sub-Tree With the Same Label

As the edges are not guaranteed to be a tree. We can either

1. rebuilt a tree
2. use a smarter way to make sure we have the tree structure.

Using the first way, got TLE as the extra time used to build the tree.

the second way can get AC.

For the whole problem, as we only have 26 alphabets. For each node, we can maintain an array which contains the occurrence of the 26 letters for each node and the subtree under that node. Then for node a, the result will be

the sum of the result from all its children node + itself.

struct pair_hash
{
 template <class T1, class T2>
 std::size_t operator () (std::pair<T1, T2> const &pair) const
 {
  std::size_t h1 = std::hash<T1>()(pair.first);
  std::size_t h2 = std::hash<T2>()(pair.second);
return h1 ^ h2;
 }
};
int node_labels[100000][26];
class Solution {
public:
    void dfs(int node, vector<vector<int>>& graph,  vector<int> &labels)
    {
        node_labels[node][labels[node]]+=1;
        //cout<<node_labels[node][labels[node]]<<endl;
        for(auto v: graph[node])
        {
            dfs(v, graph,  labels );
        }
        //node_labels[node][labels[node]-'a']+=1;
        for(auto v: graph[node])
        {
            for (int j=0;j<26;++j)
            {
                node_labels[node][j] += node_labels[v][j];
            }
        }
        return;
    }
    vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
        memset(node_labels, 0 , sizeof(node_labels));
        vector<vector<int>> graph(n);
        queue<int> q;
        q.push(0);
        map<int, vector<int>> edgess;
        unordered_set<pair<int, int>, pair_hash> pp;
        for (auto edge:edges)
        {
            pp.insert(make_pair(edge[0], edge[1]));
            pp.insert(make_pair(edge[1], edge[0]));            
            edgess[edge[0]].push_back(edge[1]);
            edgess[edge[1]].push_back(edge[0]);
        }
        while(!q.empty())
        {
            auto node = q.front();q.pop();
            for(auto v: edgess[node])
            {
                auto this_p = make_pair(v, node);
                if(pp.find(this_p)!=pp.end())
                {
                    q.push(v);
                    graph[node].push_back(v);
                    pp.erase(this_p);
                    pp.erase(make_pair(node, v));
                }
            }
        }
       
        vector<int> new_labels(n, 0);
        for(int k=0;k<n;++k)
        {
            new_labels[k] = labels[k] - 'a';
        }
       // vector<vector<int>> node_labels(n, vector<int>(26, 0));
        dfs(0, graph,  new_labels);
        vector<int> ret(n, 0);
        for(int i=0;i<n;++i)
        {
            auto this_ret = node_labels[i][new_labels[i]];
            ret[i]=this_ret;
        }
        return ret;
        
    }
};

An AC solution by using the second approach.

int node_labels[100000][26];
class Solution {
public:
    void dfs(int node, vector<vector<int>>& graph,  string &labels, vector<int>&visited)
    {
        visited[node] = 1;
        node_labels[node][labels[node]-'a']+=1;
        for(auto v: graph[node])
        {
            if(!visited[v])
            {
                dfs(v, graph,  labels, visited );
                for(int j=0;j<26;++j)
                    node_labels[node][j] += node_labels[v][j];
            }
        }
        return;
    }
    vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
        memset(node_labels, 0 , sizeof(node_labels));
        vector<vector<int>> graph(n);
        for (auto edge:edges)
        {
            graph[edge[0]].push_back(edge[1]);
            graph[edge[1]].push_back(edge[0]);
        }
        vector<int> visited(n, 0);
        dfs(0, graph,  labels, visited);
        vector<int> ret(n, 0);
        for(int i=0;i<n;++i)
        {
            auto this_ret = node_labels[i][labels[i]-'a'];
            ret[i]=this_ret;
        }
        return ret;
    }
};

Thanks for reading.