百度笔试第三题

小红拿到了一棵树,每个节点被染成了红色或者蓝色。

小红定义每条边的权值为:删除这条边时,形成的两个子树的同色连通块数量之差的绝对值。

小红想知道,所有边的权值之和是多少?

输入描述

第一行输入一个正整数 n ,代表节点的数量。

第二哈输入一个长度为 n 且仅由 R 和 B 两种字符组成的字符串。

第 i 个字符为 R 代表 i 号节点被染成红色,为 B 则被染成蓝色。

接下来的 n−1 行,每行输入两个正整数 u 和 v ,代表节点 u 和节点 v 有一条边相连。

1≤n≤200000


int main() {
    int n;
    cin >> n;
    vector<bool> colors_bl(n + 1, false);
    vector<unordered_set<int>> edges(n + 1, unordered_set<int>());
    for (int i = 1; i < n + 1; ++i) {
        char ch;
        cin >> ch;
        if (ch == 'B') {
            colors_bl[i] = true;
        }
    }
    for (int i = 0; i < n - 1; ++i) {
        int e1, e2;
        cin >> e1 >> e2;
        edges[e1].insert(e2);
        edges[e2].insert(e1);
    }
    vector<bool> is_seen(n+1, false);
    vector<int> dp(n+1, 0);
    // dfs返回结果是该节点作为根节点所包含的联通数量
    function<int(int, vector<bool>&)> dfs = [&](int i, vector<bool>& is_seen) {
        int tmp = 0;
        int sum = 0;
        vector<int> tmp_v;

        for (auto&& v : edges[i]) {
            if (is_seen[v]==false) {
                tmp++;
                is_seen[v] = true;
                int nums = dfs(v, is_seen);
                sum += nums;
                tmp_v.push_back(v);
            }
        }
        dp[i] = sum;
        if (tmp == 0) {
            dp[i] = 1;
            return 1;
        } else {
            if (tmp == 1) {
                if (colors_bl[i] != colors_bl[tmp_v[0]]) {
                    dp[i] += 1;
                }
            } else {
                if (colors_bl[i] == colors_bl[tmp_v[0]] &&
                    colors_bl[i] == colors_bl[tmp_v[1]]) {
                    dp[i] -= 1;
                } else if (colors_bl[i] != colors_bl[tmp_v[0]] &&
                           colors_bl[i] != colors_bl[tmp_v[1]]) {
                    dp[i] += 1;
                }
            }
        }
        return dp[i];
    };
    is_seen[1] = true;
    auto res = dfs(1, is_seen);
    // cout << endl;
    unordered_set<int> ss;
    int sum = 0;
    
    for (int i = 1; i <= n; ++i) {
        for (auto&& it : edges[i]) {
            int a1 = min(dp[i], dp[it]);
            int a2 = abs(res - a1);
            if (colors_bl[i] == colors_bl[it]) {
                a2++;
            } 
            sum += abs(a2 - a1);
        }
    }
    cout << sum/2 << endl;
}

#include "header.h"

int main() {
    int n;
    cin >> n;
    vector<bool> colors_bl(n + 1, false);
    vector<unordered_set<int>> edges(n + 1, unordered_set<int>());
    for (int i = 1; i < n + 1; ++i) {
        char ch;
        cin >> ch;
        if (ch == 'B') {
            colors_bl[i] = true;
        }
    }
    for (int i = 0; i < n - 1; ++i) {
        int e1, e2;
        cin >> e1 >> e2;
        edges[e1].insert(e2);
        edges[e2].insert(e1);
    }
    vector<bool> is_seen(n+1, false);
    vector<int> dp(n+1, 0);
    // dfs返回结果是该节点作为根节点所包含的联通数量
    function<int(int, vector<bool>&)> dfs = [&](int i, vector<bool>& is_seen) {
        int tmp = 0;
        int sum = 0;
        vector<int> tmp_v;

        for (auto&& v : edges[i]) {
            if (is_seen[v]==false) {
                tmp++;
                is_seen[v] = true;
                int nums = dfs(v, is_seen);
                sum += nums;
                tmp_v.push_back(v);
            }
        }
        dp[i] = sum;
        if (tmp == 0) {
            dp[i] = 1;
            return 1;
        } else {
            if (tmp == 1) {
                if (colors_bl[i] != colors_bl[tmp_v[0]]) {
                    dp[i] += 1;
                }
            } else {
                if (colors_bl[i] == colors_bl[tmp_v[0]] &&
                    colors_bl[i] == colors_bl[tmp_v[1]]) {
                    dp[i] -= 1;
                } else if (colors_bl[i] != colors_bl[tmp_v[0]] &&
                           colors_bl[i] != colors_bl[tmp_v[1]]) {
                    dp[i] += 1;
                }
            }
        }
        return dp[i];
    };
    is_seen[1] = true;
    auto res = dfs(1, is_seen);
    // cout << endl;
    unordered_set<int> ss;
    int sum = 0;
    
    for (int i = 1; i <= n; ++i) {
        for (auto&& it : edges[i]) {
            int a1 = min(dp[i], dp[it]);
            int a2 = abs(res - a1);
            if (colors_bl[i] == colors_bl[it]) {
                a2++;
            } 
            sum += abs(a2 - a1);
        }
    }
    cout << sum/2 << endl;
}
全部评论
绷不住了,百度今天笔试题和你这里的一模一样
3 回复 分享
发布于 2023-03-28 21:08 福建
多谢大佬的分享!比心惹
点赞 回复 分享
发布于 2023-03-20 12:51 内蒙古
进来狠狠学习惹
点赞 回复 分享
发布于 2023-03-20 12:57 黑龙江

相关推荐

有工作后先养猫:太好了,是超时空战警,我们有救了😋
点赞 评论 收藏
分享
7 19 评论
分享
牛客网
牛客企业服务