代码能力越来越菜了......
这题本来就是一个普普通通的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数组。