F1. Tree Cutting (Easy Version) - 题解

比赛与标签

比赛: Codeforces Round 540 (Div. 3) 标签: dfs and similar, trees 难度: *1800

题目大意喵~

主人你好呀，喵~ ฅ'ω'ฅ

这道题是关于一棵可爱的树的说！我们有一棵 n 个节点的树，每个节点可能是红色的（值为1）、蓝色的（值为2），或者是没有颜色的（值为0）。

题目要我们做的事情是，去寻找一种叫做 "nice edge"（好边）的边。如果我们把树上的一条边剪断，这棵树就会断成两个小部分（也就是两个连通块）。如果在这两个小部分中，任何一个部分都不会同时包含红色和蓝色的节点，那么我们剪断的这条边就是一条 "nice edge" 啦！

我们的任务就是数一数，这棵树里一共有多少条这样的 "nice edge" 呢？

题目还保证了，整棵树上至少有一个红色节点和至少一个蓝色节点，这点很重要哦！

解题思路大揭秘！

喵哈哈，要解决树上的问题，DFS（深度优先搜索）可是我们的超级好朋友呢！

我们的目标是判断切断一条边后，形成的两部分是否满足条件。这意味着，对于每一条边，我们都需要知道它分开的两个连通块里，各自有多少个红色节点和蓝色节点。

怎么才能高效地得到这些信息呢？思路是这样哒：

1. 预处理：首先，我们遍历一遍所有节点，数一数整棵树总共有多少个红色节点（total_red）和多少个蓝色节点（total_blue）。这就像是先清点一下我们所有的家当，喵~

2. 选择一个根并DFS：我们可以随便选一个节点当做树的根，比如说，就选1号节点吧！然后从这个根开始进行一次DFS。

3. 子树统计：在DFS的过程中，我们可以用一种“后序遍历”的方式，统计出以每个节点 u 为根的子树中，包含了多少个红色节点和多少个蓝色节点。这就像是每个节点 u 在向它的父节点汇报：“报告爸爸！我的家族（子树）里，有 x 个红孩子和 y 个蓝孩子！”。父节点收到所有孩子的汇报后，再把自己和孩子们的信息汇总起来，向上汇报给爷爷节点。

4. 遍历所有边并判断：当DFS完成后，我们就有了每个节点子树的颜色信息。现在，我们就可以遍历树上的每一条边 (p, u)（这里 p 是 u 在DFS树中的父节点）了。切断这条边，树会分成两个部分：

   • 部分一：以 u 为根的子树。这个子树里的红、蓝节点数量，我们刚刚通过DFS已经算出来啦，就是 red_count[u] 和 blue_count[u]。
   • 部分二：除了 u 的子树之外，树上剩下的所有节点。这部分的红、蓝节点数量也超级好算！就是用总数减去子树里的数量嘛，也就是 total_red - red_count[u] 和 total_blue - blue_count[u]。

5. 检验“Nice”条件：现在我们有了两部分的信息，就可以进行判断了：

   • 对于部分一（子树 u），它要满足条件，当且仅当 red_count[u] == 0 或者 blue_count[u] == 0（子树里没有红点，或者没有蓝点）。
   • 对于部分二（剩下的树），它要满足条件，当且仅当 (total_red - red_count[u]) == 0 或者 (total_blue - blue_count[u]) == 0（剩下的部分没有红点，或者没有蓝点）。

只要这两个条件同时成立，那么连接 p 和 u 的这条边就是一条可爱的 "nice edge" 啦！我们把计数器加一就好。

遍历完所有 n-1 条边，我们就得到了最终的答案！是不是感觉豁然开朗了呢，喵~？

代码实现喵~

#include <iostream>
#include <vector>

// 喵~ 这些是全局变量，用来存树的结构、颜色和统计结果的说
// 为了方便，我们用1到n的下标，和题目描述保持一致~
std::vector<int> colors;
std::vector<std::vector<int>> adj;
std::vector<int> red_count;
std::vector<int> blue_count;
int total_red = 0;
int total_blue = 0;

/**
 * @brief 进行深度优先搜索，来计算每个子树里的颜色数量
 * 
 * 对于每个节点 u，这个函数会算出以 u 为根的子树里，有多少红点点和蓝点点
 * 
 * @param u 当前节点
 * @param p 父节点（防止往回走）
 */
void dfs(int u, int p) {
    // 首先根据当前节点自己的颜色，初始化它的计数器
    if (colors[u] == 1) {
        red_count[u] = 1;
    } else if (colors[u] == 2) {
        blue_count[u] = 1;
    }

    // 递归地访问所有孩子节点，然后把它们的计数结果加到自己身上（这就是后序遍历的思想喵~）
    for (int v : adj[u]) {
        if (v == p) continue; // 不要走回头路哦！
        dfs(v, u);
        red_count[u] += red_count[v];
        blue_count[u] += blue_count[v];
    }
}

int main() {
    // 加速输入输出，让程序跑得飞快！
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(NULL);

    int n;
    std::cin >> n;

    // 根据节点数量n，调整好我们数据结构的大小
    colors.resize(n + 1);
    adj.resize(n + 1);
    red_count.assign(n + 1, 0);
    blue_count.assign(n + 1, 0);

    // 读入每个节点的颜色，并统计整棵树的红蓝总数
    for (int i = 1; i <= n; ++i) {
        std::cin >> colors[i];
        if (colors[i] == 1) {
            total_red++;
        } else if (colors[i] == 2) {
            total_blue++;
        }
    }

    // 读入n-1条边，建立邻接表来表示这棵树
    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        std::cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    // 喵~ 我们就从1号节点开始DFS吧！它的父节点可以随便设一个不存在的，比如0
    dfs(1, 0);

    int nice_edges_count = 0;
    // 接下来，我们检查每一条边是不是"nice"的
    // 这个循环很巧妙哦~ 因为我们是从1号点开始DFS的，所以从2到n的每个点 i 都有一个唯一的父节点。
    // 这样遍历 i 就等于遍历了所有 n-1 条连接子节点和父节点的边，不多也不少，喵~
    for (int i = 2; i <= n; ++i) {
        // 一条边把树分成了两块：节点i的子树，和树的其他部分。

        // 条件1: i的子树是“干净”的（没有红点或没有蓝点）
        bool subtree_ok = (red_count[i] == 0 || blue_count[i] == 0);
        
        // 计算树的其他部分有多少红蓝点
        int rest_red = total_red - red_count[i];
        int rest_blue = total_blue - blue_count[i];
        
        // 条件2: 树的其他部分也是“干净”的
        bool rest_ok = (rest_red == 0 || rest_blue == 0);

        // 如果两部分都满足条件，那这条边就是"nice"的！
        if (subtree_ok && rest_ok) {
            nice_edges_count++;
        }
    }

    std::cout << nice_edges_count << std::endl;

    return 0;
}

复杂度分析的说

• 时间复杂度: O(N) 的说 我们读入数据和建图需要 O(N) 的时间。接着进行了一次DFS，每个节点和每条边都只访问了一次，所以DFS的复杂度是 O(N)。最后，我们遍历了 N-1 条边来做判断，这也是 O(N) 的。所以总的时间复杂度就是 O(N) + O(N) + O(N) = O(N) 啦，非常高效喵！

• 空间复杂度: O(N) 的说 我们用了邻接表 adj 来存储树，它需要 O(N) 的空间。colors, red_count, blue_count 这几个数组也都是 O(N) 的大小。DFS递归时，函数调用栈的深度在最坏情况下（树是一条链）是 O(N)。所以总的空间复杂度是 O(N) 呢。

知识点与总结喵~

这道题真是一道考察树基础的可爱好题呢！通过它我们可以学到：

1. 树上DFS与子树统计: 这是解决树形问题的核心武器！通过一次DFS，我们就能预处理出所有子树的属性（比如节点数、颜色数、路径和等等），为后续的计算打下坚实的基础。
2. 换维思考: 对于一条边分割出的两个部分，我们不需要对每个部分都跑一次搜索。只需要计算出其中一个部分（子树）的信息，另一个部分的信息就可以通过 全局信息 - 子树信息 来巧妙地获得。这种思想在很多树形题目中都非常有用！
3. 问题分解: 把一个复杂的问题（判断一条边是否nice）分解成两个更简单的子问题（判断子树是否满足条件 & 判断剩余部分是否满足条件），然后组合起来，是算法解题的常见思路哦。

希望这篇题解能帮助到正在努力的你呐！继续加油，算法的世界真的很有趣，充满了各种奇思妙想！喵~ (ฅ^•ﻌ•^ฅ)