Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <bits/extc++.h>
- #define int int64_t
- using namespace std;
- using namespace __gnu_pbds;
- const int N = 1e5 + 1;
- const int mod = 1e9 + 7;
- array<bool, N> used;
- array<int, N> sz;
- array<vector<pair<int, int>>, N> g;
- int recalcSizes(int v, int p = -1) {
- sz[v] = 1;
- for (auto [u, w] : g[v]) {
- if (u != p && !used[u]) {
- sz[v] += recalcSizes(u, v);
- }
- }
- return sz[v];
- }
- int centroid(int n, int v, int p = -1) {
- for (auto [u, w] : g[v]) {
- if (u != p && !used[u] && sz[u] > n / 2) {
- return centroid(n, u, v);
- }
- }
- return v;
- }
- void dfs(int v, int p, int& ans, int &sum, int prev_sum, int& cnt, int prev_cnt, int depth) {
- cnt++;
- cnt %= mod;
- sum += depth;
- sum %= mod;
- ans += prev_cnt * depth + prev_sum;
- ans %= mod;
- for (auto [u, w] : g[v]) {
- if (u != p && !used[u]) {
- dfs(u, v, ans, sum, prev_sum, cnt, prev_cnt, depth + w);
- }
- }
- }
- int centroidDecomposition(int v) {
- recalcSizes(v);
- int ans = 0, sum = 0, cnt = 0;
- for (auto [u, w] : g[v]) {
- if (!used[u]) {
- dfs(u, v, ans, sum, sum, cnt, cnt, w);
- }
- }
- ans += sum;
- ans %= mod;
- used[v] = true;
- for (auto [u, w] : g[v]) {
- if (!used[u]) {
- ans += centroidDecomposition(centroid(sz[u], u));
- ans %= mod;
- }
- }
- return ans;
- }
- int32_t main() {
- ios::sync_with_stdio(false);
- cin.tie(nullptr);
- int n;
- cin >> n;
- for (int i = 0; i < n - 1; i++) {
- int v, u;
- cin >> v >> u;
- v--, u--;
- int w;
- cin >> w;
- g[v].emplace_back(u, w);
- g[u].emplace_back(v, w);
- }
- recalcSizes(0);
- cout << centroidDecomposition(centroid(n, 0));
- }
Advertisement
Add Comment
Please, Sign In to add comment