LeetCode "Minimum Height Tree" !!

Simple data structure but not that easy to figure out.. MHT -> balanced tree.
https://leetcode.com/problems/minimum-height-trees/

Lesson learnt: Focus on problem itself. Play with it. Do NOT jam your brain with knowledge!

class Solution 
{
public:
    vector<int> findMinHeightTrees(int n, vector<pair<int, int>>& edges) 
    {
        vector<int> ret;
        int m = edges.size();
        if(!m) return {0};
        
        //    Set up Graph
        unordered_set<int> leaves;
        unordered_set<int> all;
        for(int i = 0; i < n; i ++)
        {
            leaves.insert(i);
            all.insert(i);
        }
        if(all.size() < 3)
        {
            for(auto v: all) ret.push_back(v);
            return ret;
        }        
        
        unordered_map<int, unordered_set<int>> g;
        for(auto &p : edges)
        {
            g[p.first].insert(p.second);
            if(g[p.first].size() > 1)
                leaves.erase(p.first);
            g[p.second].insert(p.first);
            if(g[p.second].size() > 1)
                leaves.erase(p.second);
        }
        
        queue<int> q;
        for(auto l : leaves)
            q.push(l);
        
        unordered_set<int> cs;
        while(!q.empty())
        {
            int v = q.front(); q.pop();            
            all.erase(v);            
            for(auto c: g[v])
            {
                if(all.count(c))
                {
                    g[c].erase(v);
                    if(g[c].size() == 1)
                        cs.insert(c);
                }
            }
            if(q.empty())
            {
                if(all.size() <= 2) break;
            
                for(auto v : cs) q.push(v);
                cs.clear();
            }                    
        }
        for(auto v: all) ret.push_back(v);
        return ret;
    }
};