fork download
copy 
  1. #pragma GCC optimize("O3")
  2. #include <bits/stdc++.h>
  3. using namespace std;
  4.  
  5. using ll = long long;
  6.  
  7. const int MOD = 998244353;
  8. const int N = 100069;
  9. const int LOG_N = 17;
  10. const int INF = 1000000069;
  11.  
  12. int n;
  13. vector<pair<int,int>> graph[N];
  14. int weight[N];
  15.  
  16. namespace CD {
  17.     bitset<N> banned;
  18.     int sz[N];
  19.     int centroid[LOG_N][N];
  20.     int heightLevel[LOG_N][N];
  21.     int minWeight[LOG_N][N];
  22.     vector<pair<int,int>> levelList[2][N];
  23.     vector<pair<int,int>> prefixSumList[2][N];
  24.     vector<array<int,3>> prefixProdList[2][N];
  25.     int positionIndex[2][LOG_N][N];
  26.     int depthOfCentroid[N];
  27.     vector<int> distinctWeights[2][N];
  28.  
  29.     void computeSize(int u, int parent) {
  30.         sz[u] = 1;
  31.         for (auto &e : graph[u]) {
  32.             int v = e.first, id = e.second;
  33.             if (v == parent || banned[id]) continue;
  34.             computeSize(v, u);
  35.             sz[u] += sz[v];
  36.         }
  37.     }
  38.  
  39.     int findCentroid(int u, int parent, int total) {
  40.         for (auto &e : graph[u]) {
  41.             int v = e.first, id = e.second;
  42.             if (v == parent || banned[id]) continue;
  43.             if (sz[v] * 2 > total)
  44.                 return findCentroid(v, u, total);
  45.         }
  46.         return u;
  47.     }
  48.  
  49.     void dfs(int u, int parent, int layer) {
  50.         for (auto &e : graph[u]) {
  51.             int v = e.first, id = e.second;
  52.             if (v == parent || banned[id]) continue;
  53.             centroid[layer][v] = centroid[layer][u];
  54.             heightLevel[layer][v] = heightLevel[layer][u] + 1;
  55.             minWeight[layer][v] = min(minWeight[layer][u], weight[id]);
  56.             dfs(v, u, layer);
  57.         }
  58.     }
  59.  
  60.     void collect(int u, int parent, int layer, vector<pair<int,int>> &out) {
  61.         out.emplace_back(minWeight[layer][u], u);
  62.         for (auto &e : graph[u]) {
  63.             int v = e.first, id = e.second;
  64.             if (v == parent || banned[id]) continue;
  65.             collect(v, u, layer, out);
  66.         }
  67.     }
  68.  
  69.     void build(int entry, int layer) {
  70.         computeSize(entry, -1);
  71.         int c = findCentroid(entry, -1, sz[entry]);
  72.         centroid[layer][c] = c;
  73.         heightLevel[layer][c] = 0;
  74.         minWeight[layer][c] = INF;
  75.         dfs(c, -1, layer);
  76.  
  77.         vector<pair<int,int>> nodes;
  78.         collect(c, -1, layer, nodes);
  79.         sort(nodes.begin(), nodes.end());
  80.         levelList[0][c] = nodes;
  81.         depthOfCentroid[c] = layer;
  82.  
  83.         for (auto &e : graph[c]) {
  84.             int v = e.first, id = e.second;
  85.             if (banned[id]) continue;
  86.             banned[id] = true;
  87.             build(v, layer + 1);
  88.         }
  89.  
  90.         for (auto &p : nodes) {
  91.             if (p.second == c) continue;
  92.             int root = centroid[layer + 1][p.second];
  93.             levelList[1][root].push_back(p);
  94.         }
  95.     }
  96.  
  97.     inline ll square(ll x) { return (x * x) % MOD; }
  98.     inline ll modNormalize(ll x) { x %= MOD; return x < 0 ? x + MOD : x; }
  99.  
  100.     void init() {
  101.         for (int x = 0; x < 2; ++x) {
  102.             for (int i = 1; i <= n; ++i) {
  103.                 auto &vec = levelList[x][i];
  104.                 if (vec.empty()) continue;
  105.                 int m = vec.size();
  106.                 distinctWeights[x][i].resize(m);
  107.                 for (int j = 0; j < m; ++j)
  108.                     distinctWeights[x][i][j] = vec[j].first;
  109.  
  110.                 vector<int> pref(m);
  111.                 pref[0] = 0;
  112.                 for (int j = 1; j < m; ++j) {
  113.                     if (distinctWeights[x][i][j] != distinctWeights[x][i][j-1])
  114.                         pref[j] = j;
  115.                     else pref[j] = pref[j-1];
  116.                 }
  117.  
  118.                 prefixSumList[x][i].resize(m);
  119.                 prefixProdList[x][i].resize(m);
  120.                 for (int j = 0; j < m; ++j) {
  121.                     int h = heightLevel[depthOfCentroid[i] - x][ vec[j].second ];
  122.                     prefixSumList[x][i][j] = { h, square(h) };
  123.                     int w = vec[j].first;
  124.                     prefixProdList[x][i][j] = { w, (int)(1LL * h * w % MOD), (int)(1LL * square(h) * w % MOD) };
  125.                     positionIndex[x][depthOfCentroid[i]][ vec[j].second ] = pref[j];
  126.                 }
  127.  
  128.                 for (int j = m - 2; j >= 0; --j) {
  129.                     prefixSumList[x][i][j].first = (prefixSumList[x][i][j].first + prefixSumList[x][i][j+1].first) % MOD;
  130.                     prefixSumList[x][i][j].second = (prefixSumList[x][i][j].second + prefixSumList[x][i][j+1].second) % MOD;
  131.                 }
  132.                 for (int j = 1; j < m; ++j) {
  133.                     for (int k = 0; k < 3; ++k) {
  134.                         prefixProdList[x][i][j][k] = (prefixProdList[x][i][j][k] + prefixProdList[x][i][j-1][k]) % MOD;
  135.                     }
  136.                 }
  137.             }
  138.         }
  139.     }
  140.  
  141.     ll query(int u) {
  142.         ll ans = 0;
  143.         for (int level = 0; level < LOG_N; ++level) {
  144.             int c = centroid[level][u];
  145.             int idx = positionIndex[0][level][u];
  146.             ll h_u = heightLevel[level][u];
  147.             ll h_sq = square(h_u);
  148.             int m = levelList[0][c].size();
  149.             ll cur = 0;
  150.  
  151.             if (idx < m) {
  152.                 auto &S = prefixSumList[0][c][idx];
  153.                 cur = (h_sq * (m - idx) + S.second + 2 * h_u * S.first) % MOD;
  154.             }
  155.             cur = cur * minWeight[level][u] % MOD;
  156.             if (idx > 0) {
  157.                 auto &P = prefixProdList[0][c][idx-1];
  158.                 cur = (cur + h_sq * P[0] + 2 * h_u * P[1] + P[2]) % MOD;
  159.             }
  160.             ans = (ans + cur) % MOD;
  161.  
  162.             if (c == u) break;
  163.  
  164.             int c2 = centroid[level+1][u];
  165.             idx = positionIndex[1][level+1][u];
  166.             m = levelList[1][c2].size();
  167.             cur = 0;
  168.             if (idx < m) {
  169.                 auto &S = prefixSumList[1][c2][idx];
  170.                 cur = (h_sq * (m - idx) + S.second + 2 * h_u * S.first) % MOD;
  171.             }
  172.             cur = cur * minWeight[level][u] % MOD;
  173.             if (idx > 0) {
  174.                 auto &P = prefixProdList[1][c2][idx-1];
  175.                 cur = (cur + h_sq * P[0] + 2 * h_u * P[1] + P[2]) % MOD;
  176.             }
  177.             ans = (ans - cur) % MOD;
  178.             if (ans < 0) ans += MOD;
  179.         }
  180.         return ans;
  181.     }
  182. }
  183.  
  184. int main() {
  185.     ios::sync_with_stdio(false);
  186.     cin.tie(nullptr);
  187.  
  188.     if (fopen("NETW.INP", "r")) {
  189.         freopen("NETW.INP", "r", stdin);
  190.         freopen("NETW.OUT", "w", stdout);
  191.     }
  192.  
  193.     cin >> n;
  194.     for (int i = 1; i < n; ++i) {
  195.         int u, v, w;
  196.         cin >> u >> v >> w;
  197.         graph[u].push_back({v, i});
  198.         graph[v].push_back({u, i});
  199.         weight[i] = w;
  200.     }
  201.  
  202.     CD::build(1, 0);
  203.     CD::init();
  204.  
  205.     for (int i = 1; i <= n; ++i) {
  206.         cout << CD::query(i) << '\n';
  207.     }
  208.  
  209.     return 0;
  210. }
Success #stdin #stdout 0.01s 40704KB
comments (?)
stdin
copy 
7
1 2 3
1 3 2
2 4 2
2 6 1
4 5 1
4 7 2
stdout
copy 
44
26
85
35
43
36
73
https://ideone.com/rjfGB3
language:
C++14 (gcc 8.3)
created:
5 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!