leetcode

commit 0e5b14df572cb5927b6c33f0eb1c046e72684dd6
parent 649c5efa45d933d4816951d94df6022586ccc15a
Author: Dimitrije Dobrota <mail@dimitrijedobrota.com>
Date:   Sun, 28 Apr 2024 14:36:59 +0200

Daily Problem

Diffstat:
A
Problems/0834.cpp
 | 
50
++++++++++++++++++++++++++++++++++++++++++++++++++
M
README.md
 | 
1
+

2 files changed, 51 insertions(+), 0 deletions(-)
diff --git a/Problems/0834.cpp b/Problems/0834.cpp
@@ -0,0 +1,50 @@
+class Solution {
+  public:
+    vector<int> sumOfDistancesInTree(int n, vector<vector<int>> &edges) {
+        vector<vector<int>> adj(n);
+        vector<int> res(n), count(n, 1);
+
+        for (const auto &edge : edges) {
+            adj[edge[0]].push_back(edge[1]);
+            adj[edge[1]].push_back(edge[0]);
+        }
+
+        using record_t = tuple<int, int>;
+        stack<record_t> st;
+
+        for (st.emplace(-1, 0); !st.empty();) {
+            if (get<1>(st.top()) != -1) {
+                const auto [parent, root] = st.top();
+                st.emplace(-1, -1);
+                for (const auto next : adj[root]) {
+                    if (next == parent) continue;
+                    st.emplace(root, next);
+                }
+                continue;
+            }
+
+            st.pop();
+            const auto [parent, root] = st.top();
+            st.pop();
+
+            for (const auto next : adj[root]) {
+                if (next == parent) continue;
+                count[root] += count[next];
+                res[root] += res[next] + count[next];
+            }
+        }
+
+        for (st.emplace(-1, 0); !st.empty();) {
+            const auto [parent, root] = st.top();
+            st.pop();
+
+            for (const auto next : adj[root]) {
+                if (next == parent) continue;
+                res[next] = res[root] - count[next] + (n - count[next]);
+                st.emplace(root, next);
+            }
+        }
+
+        return res;
+    }
+};
diff --git a/README.md b/README.md
@@ -503,6 +503,7 @@ for solving problems.
 |  0823  |   Medium   | [Binary Trees With Factors](Problems/0823.cpp)                                                     |
 |  0830  |   Medium   | [Kth Smallest Element in a BST](Problems/0230.cpp)                                                 |
 |  0833  |   Medium   | [Find And Replace in String](Problems/0833.cpp)                                                    |
+|  0834  |    Hard    | [Sum of Distances in Tree](Problems/0834.cpp)                                                      |
 |  0835  |   Medium   | [Image Overlap](Problems/0835.cpp)                                                                 |
 |  0837  |   Medium   | [New 21 Game](Problems/0837.cpp)                                                                   |
 |  0838  |   Medium   | [Push Dominoes](Problems/0838.cpp)                                                                 |