leetcode 5465. 子树中标签相同的节点数

代码能力越来越菜了......

题目链接

这题本来就是一个普普通通的DFS，但我却掉进了坑了，一直在超时，后面看了题解以后才恍然大悟。

我一开始的想法是，dp[i][k]表示第i个节点为根的情况下，含有字符'a'+k的个数。那么从根节点开始dfs，回溯的时候累计一下就行。

我的代码是这么写的：

class Solution {
public:
    int dp[100005][26];
    bool bit[100005];
    vector<vector<int>> tree;
    vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
        memset(dp,0,sizeof(dp));
        memset(bit,0,sizeof(bit));
        tree.resize(n);
        for (auto& e : edges) 
        {
            tree[e[0]].emplace_back(e[1]);
            tree[e[1]].emplace_back(e[0]);
        }
        dfs(0,labels);
        vector<int> ans;
        for(int i=0;i<n;i++)
        {
            ans.push_back(dp[i][labels[i]-'a']);
        }
        return ans;
    }
    void dfs(int rt, string labels)
    {
        dp[rt][labels[rt]-'a']++;
        bit[rt]=1;
        for(auto vec:tree[rt])
        {
            if(bit[vec]) continue;
            dfs(vec,labels);
            for(int i=0;i<26;i++)
            {
                dp[rt][i]+=dp[vec][i];
            }
        }
    }
};

复杂度为n*26*max(节点的出度)。其实问题就出在我把for(i=0;i<26;i++)这个循环放在了遍历节点这个循环的内部。

解决方法就是利用引用传递，可以将for(i=0;i<26;i++)这个循环放在外面，从而让复杂度降低到n*max(26,max(节点的出度)）

class Solution {
    vector<vector<int>> tree;
    vector<bool> vis;
    vector<int> ans;
public:
    vector<int> countSubTrees(int n, vector<vector<int>>& edges, string labels) {
        tree.resize(labels.size());
        vis.resize(labels.size(), false);
        ans.resize(labels.size());
        for (auto& e : edges) {
            tree[e[0]].emplace_back(e[1]);
            tree[e[1]].emplace_back(e[0]);
        }
        vector<int> tmp(26);
        dfs(0, tmp, labels);
        return ans;
    }

    void dfs(int in, vector<int>& a, string& labels) {
        if (vis[in]) return; 
        vis[in] = true;
        vector<int> tmp(26);
        for (auto& out : tree[in]) dfs(out, tmp, labels);
            ans[in] = ++tmp[labels[in] - 'a'];
            for (int i = 0; i < 26; ++i) a[i] += tmp[i];
        }
};

上面代码的亮点在于处于同一层的节点共用一个tmp数组。