7. 子树边权平方和

本文主要是介绍7. 子树边权平方和，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

这是别人问我的题，所以没有题目链接，名字也是我随便起的。

有一棵包含 $n$ 个节点的树，节点编号从 $1$ 到 $n$，以 $1$ 为根，每条边都有一个权值。给你 $m$ 次询问，每次询问一个子树，在子树中对每种边权的值 $c$ 统计出现次数 $cn{t}_c$，求 $\sum (c\cdot cn{t}_c{)}^2$。

$n,m$ 以及边权范围均为 $1\times 10^5$。

容易发现可以将边权下放到点，并且这个平方和很好维护，增删一个权值都可以 $O(1)$ 完成。那么可以用 dsu on tree 来做，先预处理出所有节点的最大子树，然后 dfs 的时候先算所有小子树再算大子树，最后暴力合并小子树即可。时间复杂度 $O(n\log n)$。

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
const int maxn = 1e5 + 5;
ll weight[maxn], ans[maxn];
vector<int> g[maxn], id[maxn];
int sz[maxn], son[maxn];
void dfs0(int u) {sz[u] = 1;for (auto v : g[u]) {dfs0(v);sz[u] += sz[v];if (sz[v] > sz[son[u]])son[u] = v;}
}
int cnt[maxn], vis[maxn];
ll sum;
void clear(int u) {sum -= (weight[u] * cnt[weight[u]]) * (weight[u] * cnt[weight[u]]);cnt[weight[u]]--;sum += (weight[u] * cnt[weight[u]]) * (weight[u] * cnt[weight[u]]);for (auto v : g[u])clear(v);
}
void dfs(int u, int op) {for (auto v : g[u]) {if (v == son[u])continue;dfs(v, 0);}if (son[u]) {dfs(son[u], 1);}for (auto v : g[u]) {if (v == son[u])continue;dfs(v, 1);}if (!vis[u]) {for (auto x : id[u]) {ans[x] = sum;}vis[u] = 1;}sum -= (weight[u] * cnt[weight[u]]) * (weight[u] * cnt[weight[u]]);cnt[weight[u]]++;sum += (weight[u] * cnt[weight[u]]) * (weight[u] * cnt[weight[u]]);if (!op) {clear(u);}
}
void solve() {int n, m, c, q;cin >> n >> m >> c;for (int i = 1, u, v, w; i < n; ++i) {cin >> u >> v >> w;g[u].push_back(v);weight[v] = w;}for (int i = 1, x; i <= m; ++i) {cin >> x;id[x].push_back(i);}dfs0(1);dfs(1, 1);for (int i = 1; i <= m; ++i) {cout << ans[i] << '\n';}
}
int main() {ios::sync_with_stdio(false);cin.tie(nullptr);solve();return 0;
}

这篇关于7. 子树边权平方和的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！

---