본문으로 건너뛰기

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

info

풀이 키워드​

스포주의

dfs


풀이 코드​

info
  • 메모리: 106600 KB
  • 시간: 114 ms
class Solution {
    List<Integer>[] adj1, adj2;
    boolean[] vis;
    int n, m;
    int[] cache1; // tree1 node group per ndx

    public int dfs(int ndx, int depth, int mxDepth, int parity) {
        if (depth > mxDepth) return 0;

        int res = (depth & 1) == parity ? 1 : 0;
        if (parity == 0) cache1[ndx] = depth & 1; // mark tree1 node group
        List<Integer>[] adj = parity == 0 ? adj1 : adj2;
        for (int nxt : adj[ndx]) {
            vis[ndx] = true;
            if (!vis[nxt]) res += dfs(nxt, depth+1, mxDepth, parity);
            vis[ndx] = false;
        }

        return res;
    }

    public int[] maxTargetNodes(int[][] edges1, int[][] edges2) {
        this.n = edges1.length + 1;
        this.m = edges2.length + 1;

        // init vis
        vis = new boolean[n < m ? m : n];

        // init adj
        adj1 = new ArrayList[n];
        adj2 = new ArrayList[m];
        for (int i = 0; i < n; ++i) adj1[i] = new ArrayList<>();
        for (int j = 0; j < m; ++j) adj2[j] = new ArrayList<>();
        for (int[] e : edges1) {
            adj1[e[0]].add(e[1]);
            adj1[e[1]].add(e[0]);
        }
        for (int[] e : edges2) {
            adj2[e[0]].add(e[1]);
            adj2[e[1]].add(e[0]);
        }

        // precalc: get max from tree2
        int cand = dfs(edges2[0][0], 0, m, 1); // parity: odd
        int mx = Math.max(cand, m-cand);
        
        // precalc: get max per group from tree1
        cache1 = new int[n];
        int groupVal0 = dfs(edges1[0][0], 0, n, 0); // parity: even
        int groupVal1 = n - groupVal0;

        int[] ans = new int[n];
        for (int i = 0; i < n; ++i)
            ans[i] = (cache1[i] == 0 ? groupVal0 : groupVal1) + mx;
        
        return ans;
    }
}

풀이 해설​

3372번의 후속 문제이다.

범위가 2≤n,m≤1052 \le n, m \le 10^5 이기 때문에 dfs 호출 횟수에 유의해야 한다. (이전 문제는 10310^3이었음)


발상은 이전 문제와 비슷..​

이전 문제에선 tree1에서 k 거리, tree2에서 k-1 거리로 기준잡아 답을 구했다.

비슷한 방식으로, 이 문제에선 tree1에서 짝수 거리, tree2에서 홀수 거리로 기준잡으면 된다.


🤔 ..하지만 약간의 함정이 있음​

트리는 모든 노드가 연결되어있고, 따라서 어느 한 노드를 정했을때 짝수/홀수 거리가 되는 노드들은 일종의 그룹으로 확정된다.

예를 들어 0 - 1 - 2 - 3 - 4 형식의 트리가 있을 때, 0번 노드의 짝수/홀수 그룹은 다음과 같다.

짝수 그룹: 0, 2, 4 -> 노드 개수: 3
홀수 그룹: 1, 3 -> 노드 개수: 2

여기서 짝수 그룹의 노드 개수를 dfs 호출로 구하고, 홀수 그룹은 dfs 호출 없이 전체 노드 수에서 짝수 그룹 노드 수만큼 차감하면 된다.


코드 설명​

int cand = dfs(edges2[0][0], 0, m, 1); // parity: odd
int mx = Math.max(cand, m-cand);

tree2에선 edges2에 들어있는 맨 첫 노드를 기준으로 짝/홀 그룹을 나눈다.

사실 짝/홀이라는 표현은 어느 특정 노드를 기준잡았을때 얘기고, 그냥 그룹1 / 그룹2 라고 생각하면 된다.

연결할 tree2의 노드는 자유롭게 선택할 수 있으므로, 두 그룹 중 더 큰 카운팅 값을 mx로 저장해놓는다.


public int dfs(int ndx, int depth, int mxDepth, int parity) {
    ...
    if (parity == 0) cache1[ndx] = depth & 1; // mark tree1 node group
cache1 = new int[n];
int groupVal0 = dfs(edges1[0][0], 0, n, 0); // parity: even
int groupVal1 = n - groupVal0;

int[] ans = new int[n];
for (int i = 0; i < n; ++i)
    ans[i] = (cache1[i] == 0 ? groupVal0 : groupVal1) + mx;

tree1은 cache1이라는 캐시 배열을 만들어놓고, dfs를 돌릴때 노드별로 그룹 번호를 마킹해놓는다.

그리고 dfs 호출이 끝났을 때 반환되는 카운팅 값으로 각 그룹의 노드 수를 groupVal0, groupVal1로 저장하고,

노드를 순회할때 해당 노드의 그룹에 따라 맞는 카운팅 값을 가져다 쓰면 된다.


메모​

  • dfs 파라미터 많아진 이유: dfs1, dfs2 안 나누고 재사용하려다보니 그렇게 됨
  • 각 트리당 dfs 1번 넘게 돌리면 무조건 TLE가 터지도록 설계되어있다.
  • hard 난이도라는데, 시리즈의 이전 문제에 비해 시간복잡도가 더 까다로워졌긴 하지만 허수에 가깝다.