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


