fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. using ll = long long;
  5.  
  6. /*
  7. QUESTION IN SIMPLE WORDS
  8. ------------------------
  9. We have a rooted tree.
  10. Each node has a value.
  11.  
  12. For each node u, compute:
  13. sum of value[v] * distance(u, v)
  14. for all nodes v in subtree(u)
  15.  
  16. ------------------------------------------------------------
  17. BRUTE FORCE IDEA
  18. ------------------------------------------------------------
  19. For every node u:
  20. - iterate through all nodes in subtree(u)
  21. - compute distance(u, v)
  22. - add value[v] * distance(u, v)
  23.  
  24. ------------------------------------------------------------
  25. WHY BRUTE FORCE IS TOO SLOW
  26. ------------------------------------------------------------
  27. If we do this separately for every node,
  28. time can become O(n^2).
  29.  
  30. ------------------------------------------------------------
  31. OPTIMIZED IDEA
  32. ------------------------------------------------------------
  33. Use one DFS.
  34.  
  35. For each node u, keep:
  36. sub[u] = sum of values in subtree(u)
  37. ans[u] = required weighted sum for subtree(u)
  38.  
  39. Suppose v is a child of u.
  40.  
  41. Then every node in subtree(v) is 1 edge farther from u
  42. than it is from v.
  43.  
  44. So contribution of child subtree to ans[u] is:
  45. ans[v] + sub[v]
  46. */
  47.  
  48. struct TreeImpact {
  49.     int n;
  50.     vector<ll> a;
  51.     vector<ll> sub;
  52.     vector<ll> ans;
  53.     vector<vector<int>> g;
  54.  
  55.     TreeImpact(int n, const vector<ll>& a) : n(n), a(a) {
  56.         g.assign(n + 1, {});
  57.         sub.assign(n + 1, 0);
  58.         ans.assign(n + 1, 0);
  59.     }
  60.  
  61.     void addEdge(int u, int v) {
  62.         g[u].push_back(v);
  63.         g[v].push_back(u);
  64.     }
  65.  
  66.     void dfs(int u, int p) {
  67.         // Start with u itself
  68.         sub[u] = a[u];
  69.         ans[u] = 0;
  70.  
  71.         for (int v : g[u]) {
  72.             if (v == p) continue;
  73.  
  74.             dfs(v, u);
  75.  
  76.             // Add child subtree value
  77.             sub[u] += sub[v];
  78.  
  79.             // All nodes in child subtree become 1 farther from u
  80.             ans[u] += ans[v] + sub[v];
  81.         }
  82.     }
  83.  
  84.     vector<ll> solve(int root = 1) {
  85.         dfs(root, 0);
  86.         return ans;
  87.     }
  88. };
  89.  
  90. int main() {
  91.     int n = 5;
  92.     vector<ll> a = {0, 3, 2, 5, 10, 7};
  93.  
  94.     TreeImpact solver(n, a);
  95.     solver.addEdge(1, 2);
  96.     solver.addEdge(1, 3);
  97.     solver.addEdge(2, 4);
  98.     solver.addEdge(2, 5);
  99.  
  100.     auto ans = solver.solve(1);
  101.  
  102.     for (int i = 1; i <= n; i++) {
  103.         cout << ans[i] << ' ';
  104.     }
  105.     cout << '\n';
  106.  
  107.     return 0;
  108. }
Success #stdin #stdout 0.01s 5308KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
41 17 0 0 0
https://ideone.com/Xvu5Sl
language:
C++14 (gcc 8.3)
created:
4 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!