fork download
copy 
  1. /*
  2.     Problem: Minimum Stones to Balance a Tree
  3.     Company Tag: Tower Research OA
  4.  
  5.     Problem Statement:
  6.         We are given a tree with n nodes.
  7.  
  8.         stones[i] means initial stones on node i + 1.
  9.  
  10.         In one operation:
  11.             We can add 1 stone to any node.
  12.  
  13.         We cannot remove stones.
  14.  
  15.         The tree is balanced if for every edge (u, v):
  16.             abs(final[u] - final[v]) <= 1
  17.  
  18.         Also:
  19.             final[u] >= stones[u]
  20.  
  21.         Return minimum total stones added.
  22.  
  23.     Constraints:
  24.         2 <= n <= 200000
  25.         tree_from.length == tree_to.length == n - 1
  26.         1 <= tree_from[i], tree_to[i] <= n
  27.         0 <= stones[i] <= 1000000000
  28.         Input forms a valid tree.
  29.  
  30.     Brute Force:
  31.         Try possible final values for every node.
  32.  
  33.     Why Brute Force Fails:
  34.         n can be 200000.
  35.         stones[i] can be up to 1e9.
  36.  
  37.     Optimized Idea:
  38.         For every node u:
  39.             final[u] must be large enough because of every other node v.
  40.  
  41.         If node v has stones[v]:
  42.             final[u] >= stones[v] - distance(u, v)
  43.  
  44.         So:
  45.             final[u] = max over all v of stones[v] - distance(u, v)
  46.  
  47.         This can be computed using rerooting DP.
  48.  
  49.     Explanation Block:
  50.         down[u]:
  51.             Best requirement coming from subtree of u.
  52.  
  53.         fromParent[u]:
  54.             Best requirement coming from parent side and sibling subtrees.
  55.  
  56.         final[u]:
  57.             max(down[u], fromParent[u])
  58.  
  59.         Answer:
  60.             sum(final[u] - stones[u])
  61.  
  62.     Dry Run:
  63.         n = 3
  64.         edges:
  65.             1 - 2
  66.             1 - 3
  67.  
  68.         stones = [1, 1, 10]
  69.  
  70.         Node 3 has 10 stones.
  71.  
  72.         Node 1 is distance 1 from node 3:
  73.             node 1 must be at least 9.
  74.  
  75.         Node 2 is distance 2 from node 3:
  76.             node 2 must be at least 8.
  77.  
  78.         Final:
  79.             node 1 = 9
  80.             node 2 = 8
  81.             node 3 = 10
  82.  
  83.         Added stones:
  84.             8 + 7 + 0 = 15
  85.  
  86.     Time Complexity:
  87.         O(n)
  88.  
  89.     Space Complexity:
  90.         O(n)
  91. */
  92.  
  93. #include <bits/stdc++.h>
  94. using namespace std;
  95.  
  96. using ll = long long;
  97.  
  98. long long minimumAddedStones(
  99.     int n,
  100.     vector<int>& tree_from,
  101.     vector<int>& tree_to,
  102.     vector<long long>& stones
  103. ) {
  104.     vector<vector<int>> graph(n + 1);
  105.  
  106.     for (int i = 0; i < n - 1; i++) {
  107.         int u = tree_from[i];
  108.         int v = tree_to[i];
  109.  
  110.         graph[u].push_back(v);
  111.         graph[v].push_back(u);
  112.     }
  113.  
  114.     vector<int> parent(n + 1, 0);
  115.     vector<int> order;
  116.  
  117.     queue<int> q;
  118.  
  119.     // Root the tree at node 1.
  120.     q.push(1);
  121.     parent[1] = -1;
  122.  
  123.     while (!q.empty()) {
  124.         int u = q.front();
  125.         q.pop();
  126.  
  127.         order.push_back(u);
  128.  
  129.         for (int v : graph[u]) {
  130.             if (v == parent[u]) {
  131.                 continue;
  132.             }
  133.  
  134.             parent[v] = u;
  135.             q.push(v);
  136.         }
  137.     }
  138.  
  139.     vector<ll> down(n + 1, 0);
  140.  
  141.     // First pass:
  142.     // Calculate requirement coming from children/subtree.
  143.     for (int i = n - 1; i >= 0; i--) {
  144.         int u = order[i];
  145.  
  146.         down[u] = stones[u - 1];
  147.  
  148.         for (int v : graph[u]) {
  149.             if (parent[v] == u) {
  150.                 // If child v needs down[v],
  151.                 // parent u must be at least down[v] - 1.
  152.                 down[u] = max(down[u], down[v] - 1);
  153.             }
  154.         }
  155.     }
  156.  
  157.     const ll NEG = -4000000000000000000LL;
  158.  
  159.     vector<ll> fromParent(n + 1, NEG);
  160.  
  161.     // Second pass:
  162.     // Send parent-side information to children.
  163.     for (int u : order) {
  164.         ll best1 = NEG;
  165.         ll best2 = NEG;
  166.  
  167.         int bestChild = -1;
  168.  
  169.         // Store top two child contributions.
  170.         // This helps when excluding one child's own contribution.
  171.         for (int v : graph[u]) {
  172.             if (parent[v] == u) {
  173.                 ll contribution = down[v] - 1;
  174.  
  175.                 if (contribution > best1) {
  176.                     best2 = best1;
  177.                     best1 = contribution;
  178.                     bestChild = v;
  179.                 } else if (contribution > best2) {
  180.                     best2 = contribution;
  181.                 }
  182.             }
  183.         }
  184.  
  185.         for (int v : graph[u]) {
  186.             if (parent[v] != u) {
  187.                 continue;
  188.             }
  189.  
  190.             // If v itself gave the best contribution,
  191.             // use second best for sibling contribution.
  192.             ll bestSibling = (v == bestChild ? best2 : best1);
  193.  
  194.             ll bestAtUWithoutThisChild = stones[u - 1];
  195.  
  196.             // Requirement from parent side.
  197.             bestAtUWithoutThisChild = max(bestAtUWithoutThisChild, fromParent[u]);
  198.  
  199.             // Requirement from sibling subtrees.
  200.             bestAtUWithoutThisChild = max(bestAtUWithoutThisChild, bestSibling);
  201.  
  202.             // Moving from u to child v increases distance by 1,
  203.             // so requirement decreases by 1.
  204.             fromParent[v] = bestAtUWithoutThisChild - 1;
  205.         }
  206.     }
  207.  
  208.     long long ans = 0;
  209.  
  210.     for (int u = 1; u <= n; u++) {
  211.         ll finalValue = max(down[u], fromParent[u]);
  212.  
  213.         // We only add stones, so added amount is final - initial.
  214.         ans += finalValue - stones[u - 1];
  215.     }
  216.  
  217.     return ans;
  218. }
  219.  
  220. int main() {
  221.     int n = 3;
  222.  
  223.     vector<int> tree_from = {1, 1};
  224.     vector<int> tree_to = {2, 3};
  225.  
  226.     vector<long long> stones = {1, 1, 10};
  227.  
  228.     cout << minimumAddedStones(n, tree_from, tree_to, stones) << endl;
  229.  
  230.     return 0;
  231. }
Success #stdin #stdout 0s 5328KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
15
https://ideone.com/V9tZiU
language:
C++14 (gcc 8.3)
created:
1 day 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!