fork download
copy 
  1. /*
  2.     Problem: Sum of Jealousy Values in a Rooted Tree
  3.     Company Tag: Rubrik OA
  4.  
  5.     Problem Statement:
  6.         We are given an undirected tree with n nodes labeled from 1 to n.
  7.  
  8.         The tree is rooted at node 1.
  9.  
  10.         For each node u:
  11.             jealousy value of u = number of nodes v in subtree(u) such that v > u
  12.  
  13.         Return the sum of jealousy values of all nodes modulo 1e9 + 7.
  14.  
  15.     Constraints:
  16.         2 <= n <= 100000
  17.         A.length == B.length == n - 1
  18.         1 <= A[i], B[i] <= n
  19.         Input forms a valid tree.
  20.  
  21.     Brute Force:
  22.         For every node, traverse its subtree and count labels greater than it.
  23.  
  24.     Why Brute Force Fails:
  25.         For every node, subtree traversal can take O(n).
  26.         For all nodes, it becomes O(n^2).
  27.  
  28.     Optimized Idea:
  29.         Use Euler Tour.
  30.  
  31.         After rooting the tree:
  32.             subtree(u) becomes a continuous range:
  33.                 [tin[u], tout[u]]
  34.  
  35.         Process labels from n down to 1.
  36.  
  37.         When processing label u:
  38.             All labels greater than u are already inserted in Fenwick Tree.
  39.  
  40.     Explanation Block:
  41.         Fenwick Tree stores which nodes with greater labels have been seen.
  42.  
  43.         For current node u:
  44.             Query Fenwick Tree in range [tin[u], tout[u]]
  45.  
  46.         That gives:
  47.             number of greater labels inside subtree(u)
  48.  
  49.         Then insert u into Fenwick Tree for future smaller labels.
  50.  
  51.     Dry Run:
  52.         Tree:
  53.             1 - 2
  54.             2 - 3
  55.             1 - 4
  56.  
  57.         Node 1 subtree has 2, 3, 4:
  58.             contribution = 3
  59.  
  60.         Node 2 subtree has 3:
  61.             contribution = 1
  62.  
  63.         Node 3 and 4:
  64.             contribution = 0
  65.  
  66.         Total = 4
  67.  
  68.     Time Complexity:
  69.         O(n log n)
  70.  
  71.     Space Complexity:
  72.         O(n)
  73. */
  74.  
  75. #include <bits/stdc++.h>
  76. using namespace std;
  77.  
  78. const long long MOD = 1000000007LL;
  79.  
  80. class Fenwick {
  81. public:
  82.     int n;
  83.     vector<int> bit;
  84.  
  85.     Fenwick(int n) {
  86.         this->n = n;
  87.         bit.assign(n + 1, 0);
  88.     }
  89.  
  90.     void add(int index, int value) {
  91.         while (index <= n) {
  92.             bit[index] += value;
  93.             index += index & -index;
  94.         }
  95.     }
  96.  
  97.     int sum(int index) {
  98.         int result = 0;
  99.  
  100.         while (index > 0) {
  101.             result += bit[index];
  102.             index -= index & -index;
  103.         }
  104.  
  105.         return result;
  106.     }
  107.  
  108.     int rangeSum(int left, int right) {
  109.         return sum(right) - sum(left - 1);
  110.     }
  111. };
  112.  
  113. long long sumOfJealousyValues(int n, vector<int>& A, vector<int>& B) {
  114.     vector<vector<int>> graph(n + 1);
  115.  
  116.     for (int i = 0; i < n - 1; i++) {
  117.         graph[A[i]].push_back(B[i]);
  118.         graph[B[i]].push_back(A[i]);
  119.     }
  120.  
  121.     vector<int> parent(n + 1, -1);
  122.     vector<int> tin(n + 1);
  123.     vector<int> tout(n + 1);
  124.     vector<int> idx(n + 1, 0);
  125.  
  126.     int timer = 0;
  127.  
  128.     vector<int> st;
  129.  
  130.     // Iterative DFS to avoid recursion depth issues.
  131.     st.push_back(1);
  132.     parent[1] = 0;
  133.  
  134.     while (!st.empty()) {
  135.         int u = st.back();
  136.  
  137.         // First time visiting node u.
  138.         if (idx[u] == 0) {
  139.             tin[u] = ++timer;
  140.         }
  141.  
  142.         if (idx[u] < (int)graph[u].size()) {
  143.             int v = graph[u][idx[u]];
  144.  
  145.             idx[u]++;
  146.  
  147.             if (v == parent[u]) {
  148.                 continue;
  149.             }
  150.  
  151.             parent[v] = u;
  152.             st.push_back(v);
  153.         } else {
  154.             // All children are processed, so subtree of u ends here.
  155.             tout[u] = timer;
  156.             st.pop_back();
  157.         }
  158.     }
  159.  
  160.     Fenwick fw(n);
  161.  
  162.     long long ans = 0;
  163.  
  164.     // Process labels from larger to smaller.
  165.     for (int label = n; label >= 1; label--) {
  166.         int u = label;
  167.  
  168.         // Fenwick currently contains nodes having labels greater than u.
  169.         int greaterInsideSubtree = fw.rangeSum(tin[u], tout[u]);
  170.  
  171.         ans = (ans + greaterInsideSubtree) % MOD;
  172.  
  173.         // Insert current node for future smaller labels.
  174.         fw.add(tin[u], 1);
  175.     }
  176.  
  177.     return ans;
  178. }
  179.  
  180. int main() {
  181.     int n = 4;
  182.  
  183.     vector<int> A = {1, 2, 1};
  184.     vector<int> B = {2, 3, 4};
  185.  
  186.     cout << sumOfJealousyValues(n, A, B) << endl;
  187.  
  188.     return 0;
  189. }
Success #stdin #stdout 0s 5320KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
4
https://ideone.com/UgCd4B
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!