[LeetCode] 3203. Find Minimum Diameter After Merging Two Trees

https://leetcode.com/problems/find-minimum-diameter-after-merging-two-trees/description/

leetcode - Find Minimum Diameter After Merging Two Trees

문제 유형: 트리, bfs, 그리디

문제 난이도: Hard

 

문제

There exist two undirected trees with n and m nodes, numbered from 0 to n - 1 and from 0 to m - 1, respectively. You are given two 2D integer arrays edges1 and edges2 of lengths n - 1 and m - 1, respectively, where edges1[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the first tree and edges2[i] = [ui, vi] indicates that there is an edge between nodes ui and vi in the second tree.

You must connect one node from the first tree with another node from the second tree with an edge.

Return the minimum possible diameter of the resulting tree.

The diameter of a tree is the length of the longest path between any two nodes in the tree.

 

풀이

트리를 병합해서 최소의 지름을 만드려면, 두 루트간 잇는 것이 최선이다.

그래서, 세 가지 경우의 수가 있다:

트리 T1의 기존 지름이 병합 후 새로 생긴 지름보다 큰 경우

트리 T2의 기존 지름이 병합 후 새로 생긴 지름보다 큰 경우

새로 생긴 경로가 지름이 되는경우

 

각 트리의 지름만 구하면 세 가지 모두 알 수 있다.

트리의 지름은 어느 임의의점에서 bfs를 해서 가장 멀리 있는 한 점에서 bfs를 돌려서 가장 멀리 있는 점까지의 거리이다.

세 가지 중에서 가장 큰 값이 트리의 지름이 된다. 반환하면 된다.

 

코드

C++

#pragma GCC optimize("O3","unroll-loops");
static const int __ = [](){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    return 0;
}();

class Solution {
private:
    void generateGraph(vector<vector<int>>& edges, vector<vector<int>> &graph) {
        for(const auto &edge : edges) {
            int u = edge[0];
            int v = edge[1];
            graph[u].push_back(v);
            graph[v].push_back(u);
        }
    }

    int diameter(vector<vector<int>>& graph, int nodes) {
        queue<int> q;
        q.push(0);
        int last_visited = 0;
        vector<bool> visited(nodes, false);

        while(!q.empty()) {
            int curr = q.front();
            q.pop();

            for(int next : graph[curr]) {
                if(!visited[next]) {
                    visited[next] = true;
                    last_visited = next;
                    q.push(next);
                }
            }
        }

        q.push(last_visited);        
        fill(visited.begin(), visited.end(), false);
        visited[last_visited] = true;

        while(!q.empty()) {
            int curr = q.front();
            q.pop();

            for(int next : graph[curr]) {
                if(!visited[next]) {
                    visited[next] = true;
                    q.push(next);
                    last_visited = next;
                }
            }
        }

        q.push(last_visited);        
        fill(visited.begin(), visited.end(), false);
        vector<int> dist_factor(nodes, 0);
        visited[last_visited] = true;

        while(!q.empty()) {
            int curr = q.front();
            q.pop();

            for(int next : graph[curr]) {
                if(!visited[next]) {
                    dist_factor[next] = dist_factor[curr] + 1;
                    visited[next] = true;
                    q.push(next);
                }
            }
        }

        int diameter = *max_element(dist_factor.begin(), dist_factor.end());
        return diameter;
    }
public:
    int minimumDiameterAfterMerge(vector<vector<int>>& edges1, vector<vector<int>>& edges2) {
        int n = edges1.size()+1;
        int m = edges2.size()+1;
        vector<vector<int>> graph1(n);
        vector<vector<int>> graph2(m);
        generateGraph(edges1, graph1);
        generateGraph(edges2, graph2);
        
        int d1 = diameter(graph1, n);
        int d2 = diameter(graph2, m);

        return max({d1, d2, (d1+1)/2+(d2+1)/2+1});
    }
};