[LeetCode] 3373. Maximize the Number of Target Nodes After Connecting Trees II

https://leetcode.com/problems/maximize-the-number-of-target-nodes-after-connecting-trees-ii/description/?envType=daily-question&envId=2025-05-29

leetcode - Maximize the Number of Target Nodes After Connecting Trees II

문제 유형: bfs, dfs, 이분 그래프, 그래프, 코드

문제 난이도: Hard

 

문제

There exist two undirected trees with n and m nodes, labeled from [0, n - 1] and [0, 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.

Node u is target to node v if the number of edges on the path from u to v is even. Note that a node is always target to itself.

Return an array of n integers answer, where answer[i] is the maximum possible number of nodes that are target to node i of the first tree if you had to connect one node from the first tree to another node in the second tree.

Note that queries are independent from each other. That is, for every query you will remove the added edge before proceeding to the next query.

 

두 개의 무방향 트리가 있다. 각각 n개의 노드와 m개의 노드이다.

각 노드들은 0~n-1, 0~m-1로 라벨이 되어있다.

2D정수배열 edges1, edges2가 주어지는데, 이는 각각 1번트리와 2번트리의 간선정보가 들어있다.

 

node u에서 짝수 번 이동해서 도착한다면, node v를 타겟한다고 한다.

answer[i]에는 트리 1의 i번째 노드에서 트리 2의 어느 노드로 간선을 잇고 나서, 타겟하는 노드의 최대개수를 담으면 된다.

각 쿼리는 독립적이다.

 

풀이

트리는 이분 그래프로 표현될 수 있다.

이분 그래프로 나타내면, 서로 인접한 노드들은 서로 다른 그룹으로 있게 되고, 이것은 결국 거리의 홀짝을 구분시켜준다.

같은 색깔들은 짝수거리에 있다는 뜻이다.

 

각각 이분그래프를 통해서, 각 노드가 어느 소속이고, 각 소속의 멤버개수는 얼마인지 구해준다.

트리 2에 대해서는, 연결을 어떻게 하느냐에 따라서, 두 그룹 중 색을 마음대로 고를 수 있다.

즉, 상대 그룹은 색깔에 상관없이 더 큰걸 고르면 된다.

 

그러면, 각 쿼리에 대한 정답은:

ans[i] = (0번색이면 0번색의 트리 1에서의 개수, 1번색이면 1번색의 트리 1에서의 개수) + 연결시 최대그룹이 된다.

 

코드

C++

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

class Solution {
public:
    vector<int> getColorTable(vector<vector<int>>& edges) {
        int n = edges.size()+1;
        vector<vector<int>> adj(n);
        vector<int> color(n, -1);

        for(const auto &edge : edges) {
            int u = edge[0];
            int v = edge[1];
            adj[u].emplace_back(v);
            adj[v].emplace_back(u);
        }

        queue<int> q;
        q.emplace(0);
        color[0] = 0;
        while(!q.empty()) {
            int curr = q.front();
            q.pop();
            int currColor = color[curr];

            for(const auto &next : adj[curr]) {
                if(color[next] == -1) {
                    color[next] = currColor ^ 1;
                    q.emplace(next);
                }
            }
        }

        return color;
    }
    vector<int> maxTargetNodes(vector<vector<int>>& edges1, vector<vector<int>>& edges2) {
        int n = edges1.size()+1;
        vector<int> color1 = getColorTable(edges1);
        vector<int> color2 = getColorTable(edges2);

        int white1 = count(color1.begin(), color1.end(), 0);
        int black1 = count(color1.begin(), color1.end(), 1);
        int white2 = count(color2.begin(), color2.end(), 0);
        int black2 = count(color2.begin(), color2.end(), 1);
        int max2 = max(white2, black2);

        vector<int> ans(n);
        for(int i = 0; i < n; i++) {
            if(color1[i] == 0) ans[i] = white1 + max2;
            else ans[i] = black1 + max2;
        }

        return ans;
    }
};

 

Go

func getColors(edges [][]int) []int {
    n := len(edges) + 1
    adj := make([][]int, n)
    colors := make([]int, n)
    for i := 0; i < n; i++ {
        colors[i] = -1
    }
    for _, edge := range edges {
        u := edge[0]
        v := edge[1]
        adj[u] = append(adj[u], v)
        adj[v] = append(adj[v], u)
    }

    q := []int{0}
    colors[0] = 0
    for len(q) > 0 {
        curr := q[0]
        q = q[1:]
        for _, next := range adj[curr] {
            if colors[next] == -1 {
                colors[next] = colors[curr] ^ 1
                q = append(q, next)
            }
        }
    }

    return colors
}

func maxTargetNodes(edges1 [][]int, edges2 [][]int) []int {
    n := len(edges1)+1
    color1 := getColors(edges1)
    color2 := getColors(edges2)
    white1, black1, white2, black2 := 0, 0, 0, 0

    for _, col := range color1 {
        if col == 0 {
            white1++
        } else {
            black1++
        }
    }
    for _, col := range color2 {
        if col == 0 {
            white2++
        } else {
            black2++
        }
    }
    max2 := max(white2, black2)

    ans := make([]int, n)
    for i := 0; i < n; i++ {
        if color1[i] == 0 {
            ans[i] = white1 + max2
        } else {
            ans[i] = black1 + max2
        }
    }

    return ans
}

 

Python

from collections import deque
class Solution:
    def getColors(self, edges: List[List[int]]) -> List[int]:
        n = len(edges)+1
        adj = [[] for _ in range(n)]
        colors = [-1] * n
        for edge in edges:
            adj[edge[0]].append(edge[1])
            adj[edge[1]].append(edge[0])
        q = deque()
        q.append(0)
        colors[0] = 0

        while q:
            curr = q.popleft()
            for child in adj[curr]:
                if colors[child] == -1:
                    colors[child] = colors[curr] ^ 1
                    q.append(child)

        return colors

    def maxTargetNodes(self, edges1: List[List[int]], edges2: List[List[int]]) -> List[int]:
        n = len(edges1)+1
        color1 = self.getColors(edges1)
        color2 = self.getColors(edges2)

        white1, black1, white2, black2 = 0, 0, 0, 0
        for color in color1:
            if color == 0:
                white1 += 1
            else:
                black1 += 1
        for color in color2:
            if color == 0:
                white2 += 1
            else:
                black2 += 1
        
        max2 = max(white2, black2)

        ans = [0] * n
        for i in range(n):
            ans[i] = white1 if color1[i] == 0 else black1
            ans[i] += max2

        return ans