[M模拟] lc3249. 统计好节点的数目(dfs+图遍历+统计子树节点个数+周赛410_2)

1. 题目来源

链接：3249. 统计好节点的数目

2. 题目解析

挺不错的题目，思路本身并不难，但是比赛中的写法一直被最后一个样例卡，且比赛完后，还是被最后一个样例卡…

如果用 lambda 表达式来写 dfs 函数的话，就不会有这个问题…可能是 1e5 的情况下，传参太耗时间了吧，没搞懂这个问题。

---

正解：

• 建立无向图
• 从 0 为根开始遍历子树。
• dfs 统计子树中的节点个数。
• 判断以当前为根的各个子树中不同的节点个数，如果不同的节点个数小于等于 1 个，则统计答案。

坑点：

• 无向图统计子树节点个数时，需要记录一个 fa，然后去判断 u!=fa，才进行统计，常用的做。
• 一开始进来的时候应该将 son[u] = 1; 设置一下，因为是从根节点开始遍历的，这个无向图不会被重复遍历。

---

• 时间复杂度：$O(n)$
• 空间复杂度：$O(1)$

---

正解：

class Solution {
public:
    int n;
    unordered_map<int, vector<int>> g;
    vector<int> son;
    int res = 0;
    void dfs(int u, int fa) {
        son[u] = 1;
        set<int> st;
        for (int x : g[u]) {
            if (x == fa) continue;      // x 作为 u 的出边，不能又指回 fa，统计子树的节点个数
            dfs(x, u);                  // 递归统计
            son[u] += son[x];           // 累计子树的节点个数
            st.insert(son[x]);          // u 是根，x 是子树的各个根节点，加入 set 集合中
        }
        res += st.size() <= 1;          // 如果集合大小小于等于 1，说明没有子树或者只有一个子树，则为答案。
    }
    int countGoodNodes(vector<vector<int>>& edges) {
        n = edges.size() + 1;
        son.resize(n);
        res = 0;
        for (auto e : edges) {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }

        dfs(0, -1);

        return res;
    }
};

lambal 表达式写 dfs：

class Solution {
public:
    int countGoodNodes(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        vector<int> son(n);
        unordered_map<int, vector<int>> g;
        for (auto e : edges) {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }

        int res = 0;
        function<void(int, int)> dfs = [&](int u, int fa) {
            son[u] = 1;
            set<int> st;
            for (int x : g[u]) {
                if (x == fa) continue;
                dfs(x, u);
                son[u] += son[x];
                st.insert(son[x]);
            }
            res += st.size() <= 1;
        };

        dfs(0, -1);
        return res;
    }
};

被卡最后一个用例：

const int N = 1e5+5;
int n;
int son[N];
int f[N];
unordered_map<int, vector<int>> g;

class Solution {
public:
    void dfs(int x, int fa) {
        son[x] = 1;
        f[x] = fa;
        for (auto y : g[x]) {
            if (y == fa) continue;
            dfs(y, x);
            son[x] += son[y];
        }
    }
   
    int countGoodNodes(vector<vector<int>>& edges) {
        g = unordered_map<int, vector<int>>();
        n = edges.size() + 1;
        for (auto e : edges) {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }
        dfs(0 ,-1);

        int res = 0;
        for (int i = 0; i < n; i ++ ) {
            unordered_set<int> s;
            for (int x : g[i]) {
                if (x == f[i]) continue;
                s.insert(son[x]);
            }
            if (s.size() <= 1) res ++ ;
        }
        return res;
    }
};