Create minimum-height-trees.cpp

kamyu104 · kamyu104 · commit 259ba972bc85 · 2015-11-26T15:20:15.000+08:00
diff --git a/C++/minimum-height-trees.cpp b/C++/minimum-height-trees.cpp
@@ -0,0 +1,49 @@
+// Time:  O(n)
+// Space: O(n)
+
+class Solution {
+public:
+    vector<int> findMinHeightTrees(int n, vector<pair<int, int>>& edges) {
+        if (n == 1) {
+            return {0};
+        }
+
+        unordered_map<int, unordered_set<int>> neighbors;
+        for (const auto& e : edges) {
+            int u, v;
+            tie(u, v) = e;
+            neighbors[u].emplace(v);
+            neighbors[v].emplace(u);
+        }
+
+        vector<int> pre_level, cur_level;
+        unordered_set<int> unvisited;
+        for (int i = 0; i < n; ++i) {
+            if (neighbors[i].size() == 1) {  // A leaf.
+                pre_level.emplace_back(i);
+            }
+            unvisited.emplace(i);
+        }
+
+        // A graph can have 2 MHTs at most.
+        // BFS from the leaves until the number 
+        // of the unvisited nodes is less than 3.
+        while (unvisited.size() > 2) {
+            cur_level.clear();
+            for (const auto& u : pre_level) {
+                unvisited.erase(u);
+                for (const auto& v : neighbors[u]) {
+                    if (unvisited.count(v)) { 
+                        neighbors[v].erase(u);
+                        if (neighbors[v].size() == 1) {
+                            cur_level.emplace_back(v);
+                        }
+                    }
+                }
+            }
+            swap(pre_level, cur_level);
+        }
+        vector<int> res(unvisited.begin(), unvisited.end());
+        return res;
+    }
+};
