chheem_tapak_tam_tam

InMobi_OA_2

Oct 27th, 2025 (edited)
201
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 6.93 KB | None | 0 0
  1. #include <bits/stdc++.h> // Includes all standard libraries (common in competitive programming)
  2.  
  3. using namespace std;
  4.  
  5. // Using long long for integer type for larger values
  6. #define int long long
  7.  
  8. // Global timer for DFS start/end times
  9. int dfs_timer;
  10.  
  11. // Max log for binary lifting (2^23 > 8*10^6, safe for n up to 10^6)
  12. const int MAX_LOG = 24;
  13.  
  14. /**
  15.  * @brief Performs DFS to calculate start/end times and immediate parents.
  16.  * @param u Current node
  17.  * @param p Parent node
  18.  * @param st Start time array
  19.  * @param en End time array
  20.  * @param up Ancestor table (up[u][0] is parent of u)
  21.  * @param adj Adjacency list
  22.  */
  23. void precompute_dfs(int u, int p, vector<int>& st, vector<int>& en,
  24.                     vector<vector<int>>& up, const vector<vector<int>>& adj) {
  25.     up[u][0] = p;
  26.     st[u] = dfs_timer++;
  27.     for (int v : adj[u]) {
  28.         if (v != p) {
  29.             precompute_dfs(v, u, st, en, up, adj);
  30.         }
  31.     }
  32.     en[u] = dfs_timer;
  33. }
  34.  
  35. /**
  36.  * @brief Checks if node u is an ancestor of node v.
  37.  */
  38. bool is_ancestor(int u, int v, const vector<int>& st, const vector<int>& en) {
  39.     return (st[u] <= st[v] && en[u] >= en[v]);
  40. }
  41.  
  42. /**
  43.  * @brief Finds the Lowest Common Ancestor (LCA) of nodes u and v.
  44.  */
  45. int find_lca(int u, int v, const vector<vector<int>>& up,
  46.              const vector<int>& st, const vector<int>& en) {
  47.     // First, check if one is an ancestor of the other
  48.     if (is_ancestor(u, v, st, en)) return u;
  49.     if (is_ancestor(v, u, st, en)) return v;
  50.  
  51.     // Binary lifting to find LCA
  52.     for (int i = MAX_LOG - 1; i >= 0; i--) {
  53.         if (!is_ancestor(up[u][i], v, st, en)) {
  54.             u = up[u][i];
  55.         }
  56.     }
  57.     return up[u][0];
  58. }
  59.  
  60. signed main() {
  61.     // Optimize C++ streams for speed
  62.     ios::sync_with_stdio(false);
  63.     cin.tie(NULL);
  64.  
  65.     int n, m;
  66.     cin >> n >> m;
  67.  
  68.     vector<vector<int>> adj(n + 1);
  69.     vector<vector<int>> up(n + 1, vector<int>(MAX_LOG, 0));
  70.     vector<int> degree(n + 1, 0);
  71.     vector<int> startTime(n + 1), endTime(n + 1);
  72.  
  73.     // Read tree edges
  74.     for (int i = 0; i < n - 1; i++) {
  75.         int u, v;
  76.         cin >> u >> v;
  77.         adj[u].push_back(v);
  78.         adj[v].push_back(u);
  79.         degree[u]++;
  80.         degree[v]++;
  81.     }
  82.  
  83.     // Run DFS and precomputation, rooting at 1
  84.     // Node 0 is a virtual parent for the root
  85.     dfs_timer = 0;
  86.     precompute_dfs(1, 0, startTime, endTime, up, adj);
  87.     endTime[0] = dfs_timer; // Set end time for virtual parent
  88.  
  89.     // Fill the binary lifting (ancestor) table
  90.     for (int i = 1; i < MAX_LOG; i++) {
  91.         for (int j = 1; j <= n; j++) {
  92.             up[j][i] = up[up[j][i - 1]][i - 1];
  93.         }
  94.     }
  95.  
  96.     // 'delta' array for difference array on tree
  97.     vector<int> delta(n + 1, 0);
  98.  
  99.     // Process m queries
  100.     for (int j = 0; j < m; j++) {
  101.         int u, v;
  102.         cin >> u >> v;
  103.  
  104.         int lca = find_lca(u, v, up, startTime, endTime);
  105.  
  106.         // This is a standard difference array technique for path updates on a tree.
  107.         // Add 1 to u, add 1 to v, subtract 1 from LCA, subtract 1 from parent of LCA.
  108.         delta[u]++;
  109.         delta[v]++;
  110.         delta[lca]--;
  111.         if (up[lca][0] != 0) { // Avoid decrementing virtual parent
  112.             delta[up[lca][0]]--;
  113.         }
  114.     }
  115.  
  116.     // Use a queue to process nodes from leaves up (topological sort)
  117.     // This accumulates the delta values from children to parents
  118.     queue<int> q;
  119.     vector<int> processing_degree = degree; // Copy degrees for processing
  120.     for (int i = 1; i <= n; i++) {
  121.         if (processing_degree[i] == 1) {
  122.             q.push(i);
  123.         }
  124.     }
  125.  
  126.     // If n=1, root is also a leaf.
  127.     if (n == 1) {
  128.         q.push(1);
  129.     }
  130.    
  131.     while (!q.empty()) {
  132.         int u = q.front();
  133.         q.pop();
  134.  
  135.         int p = up[u][0];
  136.         if (p == 0) continue; // Stop at the root's virtual parent
  137.  
  138.         // Add child's sum to its parent
  139.         delta[p] += delta[u];
  140.         processing_degree[p]--;
  141.  
  142.         // If parent is now a "leaf" (relative to its own parent), add to queue
  143.         if (processing_degree[p] == 1 && p != 1) { // p != 1 ensures root isn't re-added
  144.              q.push(p);
  145.         }
  146.     }
  147.  
  148.     // After the BFS, delta[1] (the root) needs its final processing
  149.     // if it wasn't a leaf.
  150.     // The original logic handles the root correctly when its degree
  151.     // becomes 1 (or 0 if it was the last node).
  152.     // Let's re-verify the original logic.
  153.     // The original loop is correct. When deg[p] == 1, it means all children
  154.     // of p have been processed, and only the edge to *its* parent remains.
  155.     // So we push p to be processed. This works for the root 1 as well,
  156.     // as up[1][0] is 0, and the loop will just add delta[1] to delta[0]
  157.     // and terminate.
  158.     // My previous 'processing_degree' and check for p != 1 was an
  159.     // over-complication. Reverting to the simpler, correct original logic.
  160.    
  161.     // Re-doing the leaf-processing part to match the (correct) original
  162.     queue<int> q_simple;
  163.     for(int i = 1; i <= n; i++){
  164.         // Node 1 is the root, it shouldn't be added as a leaf
  165.         // unless n=1 or it's a leaf (n=2).
  166.         if(degree[i] == 1 && i != 1) {
  167.             q_simple.push(i);
  168.         }
  169.     }
  170.     // Handle n=1 case
  171.     if (n == 1) q_simple.push(1);
  172.     // Handle n=2 case (root is a leaf)
  173.     if (n == 2 && degree[1] == 1) q_simple.push(1);
  174.    
  175.     vector<int> final_ans = delta;
  176.  
  177.     while(!q_simple.empty()){
  178.         int u = q_simple.front();
  179.         q_simple.pop();
  180.  
  181.         int p = up[u][0];
  182.         if (p == 0) continue; // Should not happen if root isn't pushed incorrectly
  183.  
  184.         final_ans[p] += final_ans[u];
  185.         degree[p]--;
  186.        
  187.         // If parent (and not root) now only has parent edge, push it.
  188.         if(degree[p] == 1 && p != 1) {
  189.             q_simple.push(p);
  190.         }
  191.     }
  192.  
  193.     // The original code's logic was simpler and more robust.
  194.     // It works because even the root (1) will be pushed when degree[1] == 1,
  195.     // (after all its children are processed), and its update will
  196.     // go to node 0, which is fine.
  197.  
  198.     // --- Using the original, correct logic ---
  199.     queue<int> q_original;
  200.     for(int i = 1; i <= n; i++){
  201.         if(degree[i] == 1) {
  202.             q_original.push(i);
  203.         }
  204.     }
  205.    
  206.     // This loop correctly propagates sums up from all leaves.
  207.     while(!q_original.empty()){
  208.         int u = q_original.front();
  209.         q_original.pop();
  210.        
  211.         int p = up[u][0];
  212.         if (p == 0) continue; // Don't process parent of root
  213.  
  214.         delta[p] += delta[u];
  215.         degree[p]--;
  216.        
  217.         if(degree[p] == 1 && p != 0) { // p!=0 is technically redundant but safe
  218.             q_original.push(p);
  219.         }
  220.     }
  221.     // --- End original logic ---
  222.  
  223.  
  224.     // Print the final counts for each node
  225.     for (int i = 1; i <= n; i++) {
  226.         cout << delta[i] << ' ';
  227.     }
  228.     cout << '\n';
  229.  
  230.     return 0;
  231. }
  232.  
Advertisement
Add Comment
Please, Sign In to add comment