Leetcode Weekly Contest 197 - Count Nodes with the same number in subtrees

LINK

Level : Medium

Problem

Given a tree (i.e. a connected, undirected graph that has no cycles) consisting of n nodes numbered from 0 to n - 1 and exactly n - 1 edges. The root of the tree is the node 0, and each node of the tree has a label which is a lower-case character given in the string labels (i.e. The node with the number i has the label labels[i]).

The edges array is given on the form edges[i] = [ai, bi], which means there is an edge between nodes ai and bi in the tree.

Return an array of size n where ans[i] is the number of nodes in the subtree of the ith node which have the same label as node i.

A subtree of a tree T is the tree consisting of a node in T and all of its descendant nodes.

Solution

DFS search on tree

class Solution {
public:
    struct Node{
        int val;
        vector<Node*> children;
        Node(int in_val): val(in_val){
        }
    };
    
    void dfs(size_t root, size_t i, vector<Node*>& tree, string& labels, vector<int>& ans){
        if(i < tree.size()){
            for (auto n : tree[i]->children){
                if(labels[n->val] == labels[root]){
                    ++ans[root];
                }
                dfs(root, n->val, tree, labels, ans);
                
            }

        }

    }

    vector<int&gt countSubTrees(int n, vector<vector<int>>& edges, string labels) {
        vector<int> ans(n,1);
        vector<Node*> tree(n);
        for(size_t i = 0; i < n; ++i){
            Node* n = new Node((int)i);
            tree[i] = n;
        }
        for(auto e : edges){
            tree[e[0]]->children.push_back(tree[e[1]]);
        }
        for(size_t i = 0; i < n; ++i){
            dfs(i, i,tree, labels, ans);
        }

        return ans;
    }
};

Time copmlexity: O(n*E)