백준 18116 (로봇 조립)

문제

18116번: 로봇 조립

명령어는 I 또는 Q가 주어진다.

I가 주어질 경우 노드 2개가 추가로 주어진다.

예를 들어 I 1 2 라면 1번 노드와 2번 노드를 연결한다는 것이다.

 

Q가 주어질 경우 노드 1개가 추가로 주어진다.

예를 들어 Q 1이라면 1번 노드가 포함된 그래프의 노드 개수를 구해주면 된다.

즉 1번과 2번이 연결되있으므로 정답은 2가 된다.

---

노드를 서로 합치는 것이기 때문에 union-find 알고리즘을 사용하였다.

연결하는 것은 쉬운데 어떻게 같은 집합의 노드 개수를 셀 수 있을까???

이것은 unoin하는 과정에서 처리가 필요하다.

private static int []parent = new int[1000001]; //부모 노드 저장
private static int[]parts = new int[1000001]; //현재 노드가 포함된 집합의 노드 개수

for(int i=1;i<=1000000;i++){
	parent[i] = i;
    parts[i]=1;
}

현재는 모두 독립이므로 parts[i] = 1이 된다.

private static int find(int u){
        if(parent[u] == u)return u;
        return parent[u] = find(parent[u]);
    }

    private static void union(int a, int b){
        a = find(a);
        b = find(b);

        if(a == b)return;
        parent[a] =b;
        parts[b]+=parts[a]; //부품의 개수도 함께 합쳐짐
        parts[a]=0; //자식 노드의 부품 개수는 0으로 초기화
    }

a와 b가 합쳐지는 과정에서 a와 b는 모두 각각 최고 부모 노드이다.

따라서 합쳐질때 부품의 개수도 함께 합쳐져야한다.

StringBuilder sb = new StringBuilder();
        for(int i=0;i<n;i++){
            st = new StringTokenizer(br.readLine());

            String s = st.nextToken();

            if(s.equals("I")){
                int a = Integer.parseInt(st.nextToken());
                int b = Integer.parseInt(st.nextToken());
                union(a, b);
            }
            else{
                int a = Integer.parseInt(st.nextToken());
                a = find(a);
                sb.append(parts[a]).append("\n");
            }
        }
        System.out.println(sb);

명령어에 따라 처리하는 부분이다.

---

전체 코드이다.

public class Main {
    private static int []parent = new int[1000001];
    private static int[]parts = new int[1000001];

    private static int cnt=0;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int n = Integer.parseInt(st.nextToken());

        for(int i=1;i<=1000000;i++){
            parent[i] = i;
            parts[i]=1;
        }

        StringBuilder sb = new StringBuilder();
        for(int i=0;i<n;i++){
            st = new StringTokenizer(br.readLine());

            String s = st.nextToken();

            if(s.equals("I")){
                int a = Integer.parseInt(st.nextToken());
                int b = Integer.parseInt(st.nextToken());
                union(a, b);
            }
            else{
                int a = Integer.parseInt(st.nextToken());
                a = find(a);
                sb.append(parts[a]).append("\n");
            }
        }
        System.out.println(sb);
    }
    private static int find(int u){
        if(parent[u] == u)return u;
        return parent[u] = find(parent[u]);
    }

    private static void union(int a, int b){
        a = find(a);
        b = find(b);

        if(a == b)return;
        parent[a] =b;
        parts[b]+=parts[a];
        parts[a]=0;
    }
}