(백준) 로봇 조립 _ 18116 [java]

https://www.acmicpc.net/problem/18116

 

🐢 설명

대놓고 Union - Find 인 문제이다. 단 Q가 왔을 때 해당 부품의 로봇이 가지는 모든 부품의 개수를 반환해야 하므로 union시 부품의 개수도 같이 합쳐야 한다.

 

• I가 오면 a, b를 하나의 부품을 부모로 가지도록 union한다. 즉 같은 로봇의 부품이 된다. 이때 부모가 되는 부품은 자식이 된 부품의 개수를 더해준다.
• Q가 오면 해당 부품의 부모를 찾아내고, 해당 부모 값으로 counts배열을 접근해서 부품 개수를 가져온다.

 

paretns[] 배열을 만들어서 모든 부품의 부모 배열을 만든다.

counts[] 배열을 만들어서 로봇 별 총 부품의 개수를 담는다. (이때 초기 값은 1)

 

parents[parents[b]] = a로 함으로서 a가 최상단 부모가 되게 하고, count[a] += count[b]를 해서 b로봇의 부품개수를 a로봇의 부품 개수에 추가해준다.

 

🐢코드

package Gold이상문제_정리;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;

public class 로봇조립_18116 {
    static int N;
    static int[] parents = new int[(int) Math.pow(10, 6) + 1];
    static int[] counts = new int[(int) Math.pow(10, 6) + 1];
    static StringBuilder answer = new StringBuilder();

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine());

        for (int i = 0; i < parents.length; i++) {
            parents[i] = i;
            counts[i] = 1;
        }

        while (N-- > 0) {
            String[] input = br.readLine().split(" ");
            int partA;
            int partB;
            if (input.length == 3) {
                partA = Integer.parseInt(input[1]);
                partB = Integer.parseInt(input[2]);
                setFamily(partA, partB);
            } else if (input.length == 2) {
                partA = Integer.parseInt(input[1]);
                question(partA);
            }
        }
        System.out.print(answer);
    }

    private static void question(int part) {
        part = find(part);
        answer.append(counts[part]).append("\n");
    }

    private static void setFamily(int partA, int partB) {
        union(partA, partB);
    }

    private static void union(int partA, int partB) {
        partA = find(partA);
        partB = find(partB);
        if (partA == partB) return;
        counts[partA] += counts[partB];
        parents[parents[partB]] = partA;
    }

    private static int find(int p) {
        if (parents[p] == p) return p;
        return parents[p] = find(parents[p]);
    }
}

🐢 마무리