leetcode

Solution to some Leetcode problems written in C++

git clone git://git.dimitrijedobrota.com/leetcode.git

0834.cpp (1471B)

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