携程2023秋招笔试某题 (C++ 初级：DFS + DP 进阶：DFS + 图的前缀和)

直接看最下面进阶，第一个方法不好

题目描述

有一棵树，其中每个节点可能被染成红绿蓝(“rgb”)三种颜色，各自用一个字符表示。
现在要删除一条边，使得删除后两个连通块各自恰好包含三红颜色。
请给出合法的可以删除的边的数量。

输入

第一行一个正整数 n, 表示树的节点数量
第二行输入一个长度为 n 的，仅包含’r’, ‘g’, ‘b’的字符串，第 i 个字符串表示节点 i 的颜色。
接下来的 n - 1 行，每行两个数字，u 和 v， 代表两个节点之间有一条无向边。
$6\le n\le 10^5$
$1\le u,v\le n$

输出

合法边的数量

示例

输入

7
rgbrgbg
1 2
2 3
3 4
4 5
5 6
6 7

输出

1

说明

分析

树 + 染色 ==> DFS，这样想，问题其实就已经解决了，一个边一个边地删除，看是否符合，复杂度是 $O(n^2)$。

如果树的规模大起来$10^5$个节点，怎么办？铁超时。

那删除边后，能不能分别记录两个连通块的rgb情况？下次要遍历这个连通块的时候就不用再遍历一遍了，直接提取结果。

记录结果

• 用一个字符串‘rgb’记录，有那个颜色就加哪个字符，也可以用1 ，2， 4法。
• 记录的时候，以被删除的边的两端为目标点，记录对应的连通块。
• • 无向图的边，我们会在双方节点都进行记录，就会有边的两个方向。

树节点结构：

struct Node {
    char color;
    unordered_map<int, bool> v; // 边是否访问过
    unordered_map<int, bool> finished; // 边是否已经删除
    unordered_map<int, string> rgb; // 下一个节点对应的那一团连通块中, rgb的包含情况, 是一个字符串
};

vector<Node> nodes;

有一个边<1, 2>，边两端是1， 2

删除这条边：

• nodes[1].rgb[2]记录 2 端点代表的连通块，因为现在的方向是，从 1 到 2;
• nodes[2].rgb[1]记录 1 端点代表的连通块，因为现在的方向是，从 2 到 1;

只要在dfs时判断要访问的边是否被删除过，删除过，就直接提取结果，continue;， 然后访问下一条边。

最终时间复杂度依然是$O(n^2)$，但是已经减少很多访问了。

全部代码

#include 
#include 
#include 

using namespace std;

struct Node {
    char color;
    unordered_map<int, bool> v; // 边是否访问过
    unordered_map<int, bool> finished; // 边是否删除过
    unordered_map<int, string> rgb; // 下一个节点对应的那一团树中, rgb的包含情况, 是一个字符串
};

void dfs(vector<Node> &nodes, int index, bool &r, bool &g, bool &b) {
    if (nodes[index].color == 'r') r = true;
    if (nodes[index].color == 'g') g = true;
    if (nodes[index].color == 'b') b = true;
    if (r && g && b) // 已经全部包含就不用再找了
        return;
    for (auto node: nodes[index].v) {
        if (!node.second) { // 这条边, 在本次遍历中, 还没访问过, 那就访问
            node.second = true;
            nodes[node.first].v[index] = true;
            if (nodes[index].finished[node.first]) { 
                // 这条方向的路之前去除掉过, 那就已经有结果了
                // 就不用再继续遍历这个边了，直接看结果中存储的的rgb包含情况
                string str = nodes[index].rgb[node.first];
                for (int i = 0; i < str.size(); i++) {
                    if (str[i] == 'r') r = true;
                    if (str[i] == 'g') g = true;
                    if (str[i] == 'b') b = true;
                }
                continue; // 跳到下一条边
            }
            dfs(nodes, node.first, r, g, b);
        }
    }
}

int main() {
    int n;
    cin >> n;
    vector<Node> nodes(n);
    string colors;
    cin >> colors;
    for (int i = 0; i < n; i++) {
        nodes[i].color = colors[i];
    }
    for (int i = 1; i < n; i++) {
        int x, y;
        cin >> x >> y;
        x -= 1;
        y -= 1;
        nodes[x].v[y] = false;
        nodes[x].finished[y] = false;
        nodes[x].rgb[y] = "";
        nodes[y].v[x] = false;
        nodes[y].finished[x] = false;
        nodes[y].rgb[x] = "";
    }
    int ans = 0;
    for (int i = 0; i < n; i++) {
        for (auto node: nodes[i].finished) {
            if (!node.second) {
                vector<Node> newnodes(nodes.begin(), nodes.end()); 
                // 弄一个新的节点列表出来
                // 因为不好回溯
                newnodes[i].v[node.first] = true; // 这两行相当于"假"删除
                newnodes[node.first].v[i] = true;
                bool x, y, r, g, b;
                r = false;
                g = false;
                b = false;
                dfs(newnodes, i, r, g, b);
                if(r) nodes[node.first].rgb[i].push_back('r');
                if(g) nodes[node.first].rgb[i].push_back('g');
                if(b) nodes[node.first].rgb[i].push_back('b');
                x = r && g && b;
                r = false;
                g = false;
                b = false;
                dfs(newnodes, node.first, r, g, b);
                if(r) nodes[i].rgb[node.first].push_back('r');
                if(g) nodes[i].rgb[node.first].push_back('g');
                if(b) nodes[i].rgb[node.first].push_back('b');
                y = r && g && b;
                if (x && y) ans++;
                nodes[node.first].finished[i] = true; // 这两行是"真"删除
                node.second = true; 
            }
        }
    }
    cout << ans;
}

还可以更快，不是在删除的时候记录，而是在遍历的时候就记录。
这样的话dfs可以换成返回int值，用1 + 2 + 4分别代表rgb，初始值为0，代表没有这个方向的连通块的记录，否则不用访问，直接提取值。

进阶(图的前缀和)

1. 输入数据到graph的时候，就分别记录rgb三种颜色各自的总数。
2. DFS遍历一遍得到遍历序列 v。
3. 设置vector> pre(n + 1, vector(3)（一个int[n+1][3]的数组），注意，是 n + 1哦，pre[0][0] = pre[0][1] = pre[0][2] = 0;然后：
4. • pre[i].assign(pre[i-1].begin(), pre[i-1].end());
5. • 检查 graph[v[i]].color的值，并以此将对应颜色数字加一。
6. • 总的减去当前得到的某个颜色值，就是另一个连通块对应颜色的值。
7. • 如果6个值(两个 r g b)都大于零，就说明有一条边可以删除。

复杂度：:
$$O(n)$$

进阶代码

#include 
#include 
#include 

using namespace std;

void dfs(vector<unordered_map<int, bool>> &graph, int index, vector<int> &path) {
    path.push_back(index);
    for (auto toNode: graph[index]) {
        if (!toNode.second) {
            toNode.second = true;
            graph[toNode.first][index] = true;
            dfs(graph, toNode.first, path);
        }
    }
}

int main() {
    int n;
    cin >> n;
    vector<unordered_map<int, bool>> graph(n); // 图
    vector<int> path; // DFS路径
    vector<vector<int>> pre(n + 1, vector<int>(3, 0)); // 前缀和
    string colors;
    cin >> colors;
    vector<int> color_num(3, 0); // rgb颜色数量

    for (int i = 0; i < n; i++) { // 输入 颜色
        if (colors[i] == 'r') color_num[0]++;
        if (colors[i] == 'g') color_num[1]++;
        if (colors[i] == 'b') color_num[2]++;
    }
    for (int i = 1; i < n; i++) { // 输入 边
        int x, y;
        cin >> x >> y;
        graph[x-1][y-1] = false;
        graph[y-1][x-1] = false;
    }

    dfs(graph, 0, path); // 获取遍历序列

    int ans = 0;

    for (int i = 1; i < n - 1; i++) { // 遍历到倒数第二个就停, 后面都凑不齐三个节点了
        pre[i].assign(pre[i - 1].begin(), pre[i - 1].end()); // 复制上一个前缀和

        char color = colors[path[i - 1]]; // !!!注意这行!!!
        
        if (color == 'r') pre[i][0]++; // 对应颜色加一
        if (color == 'g') pre[i][1]++;
        if (color == 'b') pre[i][2]++;
        
        int r = color_num[0] - pre[i][0]; // 计算另一边连通块的各个颜色数量
        int g = color_num[1] - pre[i][1];
        int b = color_num[2] - pre[i][2];

        if (pre[i][0] && pre[i][1] && pre[i][2] && r && g && b) // 每个r g b都得大于 0
            ans++;
    }
    cout << ans;
}