fork download
copy 
  1. /*
  2. QUESTION 12: Media.Net OA — Choosing the Most Valuable Tree Root
  3.  
  4. Problem Statement:
  5. You are given a tree.
  6. Each node contains a candy value.
  7.  
  8. If you choose some node as the root, then every root-to-leaf path has a path sum.
  9.  
  10. For that chosen root:
  11. best = maximum root-to-leaf path sum
  12. worst = minimum root-to-leaf path sum
  13.  
  14. The value of the root is:
  15. best - worst
  16.  
  17. Return the maximum value among all possible roots.
  18.  
  19. Hard Constraints:
  20. 1 <= N <= 200000
  21. 1 <= candy[i] <= 1e9
  22.  
  23. Example:
  24. candies = [6, 2, 9, 1]
  25. edges:
  26. 1 - 2
  27. 2 - 3
  28. 2 - 4
  29.  
  30. Expected Output:
  31. 8
  32.  
  33. Explanation:
  34. If node 2 is chosen as root:
  35. paths are:
  36. 2 -> 1 = 8
  37. 2 -> 3 = 11
  38. 2 -> 4 = 3
  39.  
  40. best - worst = 11 - 3 = 8
  41.  
  42. Brute Force:
  43. For every node:
  44. - Treat it as root.
  45. - DFS to all leaves.
  46. - Find max and min root-to-leaf path sums.
  47.  
  48. Time Complexity:
  49. O(N^2)
  50.  
  51. Optimized Approach:
  52. Use rerooting DP.
  53.  
  54. For each directed side of an edge, compute the best and worst path from that side
  55. to a leaf in that component.
  56.  
  57. Then every node can combine values from all neighboring directions.
  58.  
  59. Time Complexity:
  60. O(N)
  61.  
  62. Space Complexity:
  63. O(N)
  64. */
  65.  
  66. #include <bits/stdc++.h>
  67. using namespace std;
  68.  
  69. class Solution {
  70. public:
  71.     using ll = long long;
  72.  
  73.     const ll NEG = -(1LL << 60);
  74.     const ll POS =  (1LL << 60);
  75.  
  76.     int n;
  77.     vector<vector<int>> g;
  78.     vector<ll> val;
  79.  
  80.     vector<ll> downMax, downMin;
  81.     ll answer = 0;
  82.  
  83.     void dfsDown(int u, int p) {
  84.         downMax[u] = NEG;
  85.         downMin[u] = POS;
  86.  
  87.         bool hasChild = false;
  88.  
  89.         for (int v : g[u]) {
  90.             if (v == p) continue;
  91.  
  92.             hasChild = true;
  93.             dfsDown(v, u);
  94.  
  95.             downMax[u] = max(downMax[u], val[u] + downMax[v]);
  96.             downMin[u] = min(downMin[u], val[u] + downMin[v]);
  97.         }
  98.  
  99.         if (!hasChild) {
  100.             downMax[u] = val[u];
  101.             downMin[u] = val[u];
  102.         }
  103.     }
  104.  
  105.     void dfsReroot(int u, int p, ll fromParentMax, ll fromParentMin) {
  106.         vector<pair<ll, int>> maxCand;
  107.         vector<pair<ll, int>> minCand;
  108.  
  109.         if (fromParentMax != NEG) {
  110.             maxCand.push_back({fromParentMax, p});
  111.             minCand.push_back({fromParentMin, p});
  112.         }
  113.  
  114.         for (int v : g[u]) {
  115.             if (v == p) continue;
  116.  
  117.             maxCand.push_back({downMax[v], v});
  118.             minCand.push_back({downMin[v], v});
  119.         }
  120.  
  121.         if (!maxCand.empty()) {
  122.             ll rootMax = val[u] + max_element(maxCand.begin(), maxCand.end())->first;
  123.             ll rootMin = val[u] + min_element(minCand.begin(), minCand.end())->first;
  124.  
  125.             answer = max(answer, rootMax - rootMin);
  126.         }
  127.  
  128.         for (int v : g[u]) {
  129.             if (v == p) continue;
  130.  
  131.             ll bestMax = NEG;
  132.             ll bestMin = POS;
  133.  
  134.             for (auto [x, id] : maxCand) {
  135.                 if (id != v) bestMax = max(bestMax, x);
  136.             }
  137.  
  138.             for (auto [x, id] : minCand) {
  139.                 if (id != v) bestMin = min(bestMin, x);
  140.             }
  141.  
  142.             ll passMax = (bestMax == NEG ? val[u] : val[u] + bestMax);
  143.             ll passMin = (bestMin == POS ? val[u] : val[u] + bestMin);
  144.  
  145.             dfsReroot(v, u, passMax, passMin);
  146.         }
  147.     }
  148.  
  149.     long long bestRootValue(vector<int>& candies, vector<vector<int>>& edges) {
  150.         n = candies.size();
  151.  
  152.         if (n == 1) return 0;
  153.  
  154.         val.assign(n, 0);
  155.  
  156.         for (int i = 0; i < n; i++) {
  157.             val[i] = candies[i];
  158.         }
  159.  
  160.         g.assign(n, {});
  161.  
  162.         for (auto &e : edges) {
  163.             int u = e[0] - 1;
  164.             int v = e[1] - 1;
  165.  
  166.             g[u].push_back(v);
  167.             g[v].push_back(u);
  168.         }
  169.  
  170.         downMax.assign(n, NEG);
  171.         downMin.assign(n, POS);
  172.  
  173.         dfsDown(0, -1);
  174.         dfsReroot(0, -1, NEG, POS);
  175.  
  176.         return answer;
  177.     }
  178. };
  179.  
  180. int main() {
  181.     Solution sol;
  182.  
  183.     vector<int> candies = {6, 2, 9, 1};
  184.  
  185.     vector<vector<int>> edges = {
  186.         {1, 2},
  187.         {2, 3},
  188.         {2, 4}
  189.     };
  190.  
  191.     cout << sol.bestRootValue(candies, edges) << endl;
  192.  
  193.     return 0;
  194. }
Success #stdin #stdout 0s 5320KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
8
https://ideone.com/Z9xgla
language:
C++14 (gcc 8.3)
created:
9 hours 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!