[백준] 18116. 로봇 조립

문제

18116번: 로봇 조립

풀이

쿼리에 맞게 집합을 형성하거나(I) 집합 내의 원소 개수를 출력(Q) 하는 문제

 

집합 형성(I): a와 b를 union 해준다.(여기서 union은 union-find의 메서드와 동일)

원소의 개수 출력(Q)을 위해서는 cnt배열(초기화 1)을 만들어 union 과정에서 값을 더해주면 된다.

출력 값은 cnt[find(찾는 값)]

코드

import java.io.*;
import java.util.*;

public class Main {
    static int N;
    static int[] parents, cnt;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        StringBuilder sb = new StringBuilder();

        N = Integer.parseInt(st.nextToken());

        parents = new int[1000001];
        cnt = new int[1000001];

        // parents 배열 초기화
        makeSet();

        for(int i = 0 ; i < N ; i ++){
            st = new StringTokenizer(br.readLine());
            String type = st.nextToken();
            int a = Integer.parseInt(st.nextToken());

            // union
            if(type.equals("I")){
                int b = Integer.parseInt(st.nextToken());
                union(a, b);
            }
            // print
            else sb.append(cnt[find(a)]).append("\n");
        }

        System.out.print(sb);
    }

    private static void makeSet() {
        for(int i = 1 ; i <= 1000000 ; i ++){
            parents[i] = i;
            cnt[i] = 1;
        }
    }

    private static void union(int x, int y) {
        x = find(x);
        y = find(y);

        if(x != y){
            cnt[x] += cnt[y];
            parents[parents[y]] = x;
        }
    }

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

시간/공간 복잡도

1,280ms / 257,772 KB