Description

Submission

class Solution {
public:
    vector<int> findMinHeightTrees(int n, vector<vector<int>>& edges) {
        if(n == 1) return {0};
        if(n == 2) return {0, 1};
        vector<vector<int>> next(n);
        vector<int> degree(n, 0);
        for(auto e: edges) {
            next[e[0]].push_back(e[1]);
            next[e[1]].push_back(e[0]);
            degree[e[0]]++;
            degree[e[1]]++;
        }

        queue<int> q;

        for(int i = 0; i < n; ++i) {
            if(degree[i] == 1) q.push(i);
        }

        int count = n;

        while(!q.empty()) {
            int len = q.size();

            while(len--) {
                int cur = q.front();
                q.pop();
                degree[cur] = 0;
                count--;

                for(auto val: next[cur]) {
                    if(degree[val] > 0) {
                        degree[val]--;
                    }
                    if(degree[val] == 1) {
                        q.push(val);
                    }
                }
            }
            if(count <= 2) break;
        }

        vector<int> ret;
        while(!q.empty()) {
            ret.push_back(q.front());
            q.pop();
        }

        return ret;
    }
};

220130 Update

class Solution {
public:
    vector<int> findMinHeightTrees(int n, vector<vector<int>>& edges) {
        if(n == 1) return {0};
        if(n == 2) return {0, 1};
        
        vector<vector<int>> graph(n);
        vector<int> degree(n);

        for(auto& edge: edges) {
            int a = edge[0], b = edge[1];
            degree[a]++;
            degree[b]++;
            graph[a].push_back(b);
            graph[b].push_back(a);
        }

        queue<int> q;
        for(int i = 0; i < n; ++i) {
            if(degree[i] == 1) q.push(i);
        }

        int t = n;

        while(!q.empty()) {
            int cnt = q.size();

            while(cnt--) {
                int cur = q.front();
                q.pop();
                --t;
                degree[cur] = 0;

                for(int x: graph[cur]) {
                    degree[x]--;
                    if(degree[x] == 1) {
                        q.push(x);
                    }
                }
            }
            if(t <= 2) {
                break;
            }
        }

        vector<int> rets;
        while(!q.empty()) {
            rets.push_back(q.front());
            q.pop();
        }
        return rets;
    }
};