leetcode

Solution to some Leetcode problems written in C++
git clone git://git.dimitrijedobrota.com/leetcode.git
Log | Files | Refs | README | LICENSE

0834.cpp (1471B)

      1 class Solution {
      2   public:
      3     vector<int> sumOfDistancesInTree(int n, vector<vector<int>> &edges) {
      4         vector<vector<int>> adj(n);
      5         vector<int> res(n), count(n, 1);
      6 
      7         for (const auto &edge : edges) {
      8             adj[edge[0]].push_back(edge[1]);
      9             adj[edge[1]].push_back(edge[0]);
     10         }
     11 
     12         using record_t = tuple<int, int>;
     13         stack<record_t> st;
     14 
     15         for (st.emplace(-1, 0); !st.empty();) {
     16             if (get<1>(st.top()) != -1) {
     17                 const auto [parent, root] = st.top();
     18                 st.emplace(-1, -1);
     19                 for (const auto next : adj[root]) {
     20                     if (next == parent) continue;
     21                     st.emplace(root, next);
     22                 }
     23                 continue;
     24             }
     25 
     26             st.pop();
     27             const auto [parent, root] = st.top();
     28             st.pop();
     29 
     30             for (const auto next : adj[root]) {
     31                 if (next == parent) continue;
     32                 count[root] += count[next];
     33                 res[root] += res[next] + count[next];
     34             }
     35         }
     36 
     37         for (st.emplace(-1, 0); !st.empty();) {
     38             const auto [parent, root] = st.top();
     39             st.pop();
     40 
     41             for (const auto next : adj[root]) {
     42                 if (next == parent) continue;
     43                 res[next] = res[root] - count[next] + (n - count[next]);
     44                 st.emplace(root, next);
     45             }
     46         }
     47 
     48         return res;
     49     }
     50 };