百度笔试第三题
小红拿到了一棵树,每个节点被染成了红色或者蓝色。
小红定义每条边的权值为:删除这条边时,形成的两个子树的同色连通块数量之差的绝对值。
小红想知道,所有边的权值之和是多少?
输入描述
第一行输入一个正整数 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; }