fork download 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. using ll = long long;
  5.  
  6. /*
  7. ==================== PROBLEM STATEMENT ====================
  8.  
  9. You are given a rooted tree with n nodes, rooted at node 1.
  10. Each node i has an integer value value[i].
  11.  
  12. Define subtreeSum[u] as the sum of all values in the subtree
  13. of node u (including u itself).
  14.  
  15. You are allowed to perform at most ONE operation:
  16. - Choose any node u
  17. - Change its value to ANY integer (i.e., add some delta)
  18.  
  19. Effect of this operation:
  20. - Only node u and all its ancestors (up to the root)
  21.   will have their subtree sums changed.
  22.  
  23. Goal:
  24. After performing at most one operation, maximize the frequency
  25. of any subtree sum value.
  26.  
  27. In other words:
  28. Make as many nodes as possible have the same subtree sum.
  29.  
  30. Output:
  31. Print the maximum possible frequency.
  32.  
  33. Constraints:
  34. - 1 ≤ n ≤ (typical large constraints)
  35. - Values can be large → subtree sums can be large
  36.  
  37. ===========================================================
  38. */
  39.  
  40. int main() {
  41.     ios::sync_with_stdio(false);
  42.     cin.tie(nullptr);
  43.  
  44.     int n;
  45.     cin >> n;
  46.  
  47.     // Read values of nodes (1-indexed)
  48.     vector<ll> value(n + 1);
  49.     for (int i = 1; i <= n; i++) {
  50.         cin >> value[i];
  51.     }
  52.  
  53.     /*
  54.     Parent input may come in two formats:
  55.     1) parent[1..n] where parent[1] = 0
  56.     2) parent[2..n]
  57.     So we read all remaining input and detect format
  58.     */
  59.     vector<int> rest;
  60.     int x;
  61.     while (cin >> x) {
  62.         rest.push_back(x);
  63.     }
  64.  
  65.     vector<int> parent(n + 1, 0);
  66.  
  67.     if ((int)rest.size() == n) {
  68.         // Full parent array given
  69.         for (int i = 1; i <= n; i++) {
  70.             parent[i] = rest[i - 1];
  71.         }
  72.     } else {
  73.         // parent[2..n] given
  74.         for (int i = 2; i <= n; i++) {
  75.             parent[i] = rest[i - 2];
  76.         }
  77.     }
  78.  
  79.     // Build adjacency list (tree)
  80.     vector<vector<int>> tree(n + 1);
  81.     for (int i = 2; i <= n; i++) {
  82.         tree[parent[i]].push_back(i);
  83.     }
  84.  
  85.     /*
  86.     STEP 1: Compute subtree sums
  87.  
  88.     We do iterative DFS to get traversal order,
  89.     then process in reverse (postorder style)
  90.     */
  91.     vector<int> order;
  92.     order.reserve(n);
  93.  
  94.     stack<int> st;
  95.     st.push(1);
  96.  
  97.     while (!st.empty()) {
  98.         int u = st.top();
  99.         st.pop();
  100.         order.push_back(u);
  101.  
  102.         for (int v : tree[u]) {
  103.             st.push(v);
  104.         }
  105.     }
  106.  
  107.     // subSum[u] = sum of subtree rooted at u
  108.     vector<ll> subSum(n + 1, 0);
  109.  
  110.     // Process in reverse order (children before parent)
  111.     for (int i = n - 1; i >= 0; i--) {
  112.         int u = order[i];
  113.         subSum[u] = value[u];
  114.  
  115.         for (int v : tree[u]) {
  116.             subSum[u] += subSum[v];
  117.         }
  118.     }
  119.  
  120.     /*
  121.     STEP 2: Coordinate compression
  122.  
  123.     Subtree sums can be large, so map them to [0..m-1]
  124.     */
  125.     vector<ll> allSums;
  126.     allSums.reserve(n);
  127.  
  128.     for (int i = 1; i <= n; i++) {
  129.         allSums.push_back(subSum[i]);
  130.     }
  131.  
  132.     sort(allSums.begin(), allSums.end());
  133.     allSums.erase(unique(allSums.begin(), allSums.end()), allSums.end());
  134.  
  135.     int m = (int)allSums.size();
  136.  
  137.     // sumId[u] = compressed id of subtree sum of node u
  138.     vector<int> sumId(n + 1);
  139.  
  140.     // totalFreq[id] = how many nodes have this subtree sum
  141.     vector<int> totalFreq(m, 0);
  142.  
  143.     for (int i = 1; i <= n; i++) {
  144.         sumId[i] = lower_bound(allSums.begin(), allSums.end(), subSum[i]) - allSums.begin();
  145.         totalFreq[sumId[i]]++;
  146.     }
  147.  
  148.     /*
  149.     STEP 3: DFS over all root-to-node paths
  150.  
  151.     pathFreq[id] = how many times this sum appears on current path
  152.     totalFreq[id] - pathFreq[id] = count outside path
  153.  
  154.     We maintain:
  155.     - onPath: multiset of frequencies on current path
  156.     - outsidePath: multiset of frequencies outside path
  157.  
  158.     This allows O(log n) retrieval of max frequency
  159.     */
  160.     vector<int> pathFreq(m, 0);
  161.     multiset<int> onPath, outsidePath;
  162.  
  163.     // Initially, path is empty → all nodes are outside
  164.     for (int i = 0; i < m; i++) {
  165.         onPath.insert(0);
  166.         outsidePath.insert(totalFreq[i]);
  167.     }
  168.  
  169.     int answer = 1;
  170.  
  171.     /*
  172.     DFS using explicit stack:
  173.     state = 0 → entering node
  174.     state = 1 → exiting node (backtracking)
  175.     */
  176.     vector<pair<int, int>> dfs;
  177.     dfs.push_back({1, 0});
  178.  
  179.     while (!dfs.empty()) {
  180.         auto [u, state] = dfs.back();
  181.         dfs.pop_back();
  182.  
  183.         int id = sumId[u];
  184.  
  185.         if (state == 0) {
  186.             // -------- ENTER NODE --------
  187.  
  188.             // Remove old values before updating
  189.             onPath.erase(onPath.find(pathFreq[id]));
  190.             outsidePath.erase(outsidePath.find(totalFreq[id] - pathFreq[id]));
  191.  
  192.             // Include this node in path
  193.             pathFreq[id]++;
  194.  
  195.             // Insert updated values
  196.             onPath.insert(pathFreq[id]);
  197.             outsidePath.insert(totalFreq[id] - pathFreq[id]);
  198.  
  199.             /*
  200.             Best result if we modify this path:
  201.             = best freq on path + best freq outside path
  202.             */
  203.             int bestOnPath = *onPath.rbegin();
  204.             int bestOutside = *outsidePath.rbegin();
  205.  
  206.             answer = max(answer, bestOnPath + bestOutside);
  207.  
  208.             // Schedule exit (backtracking)
  209.             dfs.push_back({u, 1});
  210.  
  211.             // Visit children
  212.             for (int i = (int)tree[u].size() - 1; i >= 0; i--) {
  213.                 dfs.push_back({tree[u][i], 0});
  214.             }
  215.  
  216.         } else {
  217.             // -------- EXIT NODE --------
  218.  
  219.             // Remove current values
  220.             onPath.erase(onPath.find(pathFreq[id]));
  221.             outsidePath.erase(outsidePath.find(totalFreq[id] - pathFreq[id]));
  222.  
  223.             // Remove node from path
  224.             pathFreq[id]--;
  225.  
  226.             // Insert updated values
  227.             onPath.insert(pathFreq[id]);
  228.             outsidePath.insert(totalFreq[id] - pathFreq[id]);
  229.         }
  230.     }
  231.  
  232.     cout << answer << '\n';
  233.     return 0;
  234. }
Runtime error #stdin #stdout 0.04s 5336KB
comments (?)
stdin 
Standard input is empty
stdout 
Standard output is empty
https://ideone.com/3BTbfz
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!