A22618. 给定一棵包含 N个顶点的树。顶点编号为 1 至 N,第 i 条边 (1≤ i ≤ N−1) 连接顶点 ai 与顶点 bi。对于树中任意两个顶点 u 和 v(满足 u<v),定义距离 d(u,v)为连接 u 和 v的简单路径上的边的数量。请计算所有满足 u<v 的顶点对 (u,v) 的距离 d(u,v) 的总和。
填空题
困难
知识点
题目描述
题目描述
给定一棵包含 N个顶点的树。顶点编号为 1 至 N,第 i 条边 (1≤ i ≤ N−1) 连接顶点 ai 与顶点 bi。
对于树中任意两个顶点 u 和 v(满足 u<v),定义距离 d(u,v)为连接 u 和 v的简单路径上的边的数量。
请计算所有满足 u<v 的顶点对 (u,v) 的距离 d(u,v) 的总和。
输入格式
第一行 ,一个整数表示 n
接下来的n−1行,每行两个整数ai,bi.
输出格式
输出所有满足 u<v的顶点对 (u,v) 的距离 d(u,v)的总和。
输入样例#1
3
1 2
2 3输出样例#1
3输入样例#2
5
1 2
1 3
1 4
1 5输出样例#2
10数据范围
2≤N≤105,1≤ai,bi≤N
参考答案
#include <iostream>
#include <vector>
#include <stack>
using namespace std;
typedef long long ll;
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n;
cin >> n;
// 建邻接表(节点编号1~n)
vector<vector<int>> adj(n + 1);
for (int i = 0; i < n - 1; ++i) {
int a, b;
cin >> a >> b;
adj[a].push_back(b);
adj[b].push_back(a);
}
vector<int> size(n + 1, 1); // 每个节点初始子树大小为1(自身)
vector<bool> visited(n + 1, false);
ll ans = 0;
// 迭代式DFS(后序遍历),栈中存储(节点, 父节点, 是否已处理子节点)
stack<tuple<int, int, bool>> st;
st.emplace(1, -1, false);
while (!st.empty()) {
auto [u, parent, processed] = st.top();
st.pop();
if (!processed) {
if (visited[u]) continue;
visited[u] = true;
// 先标记为未处理,重新压入栈,再压入所有子节点
st.emplace(u, parent, true);
for (int v : adj[u]) {
if (v != parent && !visited[v]) {
st.emplace(v, u, false);
}
}
} else {
// 后序处理:累加子节点的size,并计算边的贡献
for (int v : adj[u]) {
if (v != parent) {
size[u] += size[v];
ans += (ll)size[v] * (n - size[v]);
}
}
}
}
cout << ans << endl;
return 0;
}
上一题
下一题