fork download 
  1. /*******************************************************
  2. 2) Intuit OA — Count Y-Shaped Triplets in a Tree
  3.  
  4. Problem:
  5. - Count unordered triplets {a,b,c} in a tree such that their minimal connecting
  6.   structure is a "Y" with a junction node.
  7.  
  8. Idea:
  9. - Fix junction node x.
  10. - Remove x -> tree splits into components of sizes s1, s2, ...
  11. - Triplet is valid iff picks 1 node from 3 different components.
  12.   Count at x = sum_{i<j<k} si*sj*sk
  13. - Sum over all x
  14.  
  15. Time: O(n)
  16. *******************************************************/
  17.  
  18. #include <bits/stdc++.h>                 // standard CP header
  19. using namespace std;                     // use std names directly
  20.  
  21. int main() {                             // entry point
  22.     ios::sync_with_stdio(false);         // fast I/O
  23.     cin.tie(nullptr);                    // fast I/O
  24.  
  25.     int n;                               // number of nodes
  26.     cin >> n;                            // read n
  27.  
  28.     vector<vector<int>> g(n + 1);        // adjacency list (1-indexed)
  29.     for (int i = 0; i < n - 1; i++) {    // read n-1 edges
  30.         int u, v;                        // endpoints
  31.         cin >> u >> v;                   // input edge
  32.         g[u].push_back(v);               // add v to u
  33.         g[v].push_back(u);               // add u to v (tree is undirected)
  34.     }
  35.  
  36.     vector<int> parent(n + 1, 0);        // parent[u] in rooted tree
  37.     vector<int> sub(n + 1, 0);           // sub[u] = subtree size of u
  38.     parent[1] = -1;                      // root has no parent
  39.  
  40.     vector<int> order;                   // to store DFS order
  41.     order.reserve(n);                    // reserve memory to avoid reallocations
  42.     stack<int> st;                       // stack for iterative DFS
  43.     st.push(1);                          // start DFS from node 1
  44.  
  45.     while (!st.empty()) {                // while nodes to process exist
  46.         int u = st.top();                // take top
  47.         st.pop();                        // remove it
  48.         order.push_back(u);              // store visit order
  49.         for (int v : g[u]) {             // scan neighbors
  50.             if (v == parent[u]) continue;// ignore edge back to parent
  51.             parent[v] = u;               // set parent
  52.             st.push(v);                  // push child
  53.         }
  54.     }
  55.  
  56.     for (int i = (int)order.size() - 1; i >= 0; i--) { // process in reverse order (postorder)
  57.         int u = order[i];                // current node
  58.         int sz = 1;                      // count self
  59.         for (int v : g[u]) {             // scan neighbors
  60.             if (v == parent[u]) continue;// skip parent
  61.             sz += sub[v];                // add child subtree size
  62.         }
  63.         sub[u] = sz;                     // store subtree size
  64.     }
  65.  
  66.     long long total = 0;                 // final answer
  67.  
  68.     for (int x = 1; x <= n; x++) {       // try each node as junction
  69.         vector<long long> comps;         // component sizes after removing x
  70.         comps.reserve(g[x].size());      // reserve approximately degree(x)
  71.  
  72.         for (int v : g[x]) {             // for each neighbor of x
  73.             if (v == parent[x]) {        // neighbor is parent side
  74.                 comps.push_back((long long)n - sub[x]); // size of "everything outside x's subtree"
  75.             } else {                     // neighbor is child side
  76.                 comps.push_back(sub[v]); // size of child's subtree
  77.             }
  78.         }
  79.  
  80.         long long sum1 = 0;              // sum of sizes processed so far
  81.         long long sum2 = 0;              // sum of pair-products so far
  82.         long long ans_x = 0;             // triplets contributed by junction x
  83.  
  84.         for (long long s : comps) {      // for each component size
  85.             ans_x += s * sum2;           // add triples ending with this component
  86.             sum2  += s * sum1;           // update pair sums
  87.             sum1  += s;                  // update size sum
  88.         }
  89.  
  90.         total += ans_x;                  // add contribution for x
  91.     }
  92.  
  93.     cout << total << "\n";               // print answer
  94.     return 0;                            // done
  95. }
Runtime error #stdin #stdout 0.05s 5316KB
comments (?)
stdin 
Standard input is empty
stdout 
Standard output is empty
https://ideone.com/rBMXzG
language:
C++14 (gcc 8.3)
created:
117 days ago
visibility:
secret

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!