문제

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

민혁이는 소셜 네트워크 사이트에서 친구를 만드는 것을 좋아하는 친구이다. 우표를 모으는 취미가 있듯이, 민혁이는 소셜 네트워크 사이트에서 친구를 모으는 것이 취미이다.

어떤 사이트의 친구 관계가 생긴 순서대로 주어졌을 때, 두 사람의 친구 네트워크에 몇 명이 있는지 구하는 프로그램을 작성하시오.

친구 네트워크란 친구 관계만으로 이동할 수 있는 사이를 말한다.

입력
첫째 줄에 테스트 케이스의 개수가 주어진다. 각 테스트 케이스의 첫째 줄에는 친구 관계의 수 F가 주어지며, 이 값은 100,000을 넘지 않는다. 다음 F개의 줄에는 친구 관계가 생긴 순서대로 주어진다. 친구 관계는 두 사용자의 아이디로 이루어져 있으며, 알파벳 대문자 또는 소문자로만 이루어진 길이 20 이하의 문자열이다.

출력
친구 관계가 생길 때마다, 두 사람의 친구 네트워크에 몇 명이 있는지 구하는 프로그램을 작성하시오.

예제 입력 1
2
3
Fred Barney
Barney Betty
Betty Wilma
3
Fred Barney
Betty Wilma
Barney Betty
예제 출력 1
2
3
4
2
2
4

접근

친구 관계가 생길 때마다, 즉 각 줄의 입력이 들어올 때마다 두 사람이 속해 있는 네트워크가 하나로 합쳐지므로 유니온 파인드로 분리 집합을 구현하면 되는 문제이다.

다른 분리 집합 문제와 조금 달랐던 부분은 다음과 같다.
1. 노드가 번호가 아닌 문자열 형태로 주어짐
보통 유니온 파인드는 정수 배열을 자신의 인덱스 값으로 초기화한 후 원소의 부모 노드에 인덱스로 접근하는데, 이 문제는 입력이 문자열로 주어져서 입력을 그대로 인덱스로 사용할 수가 없었다. 그래서 HashMap을 사용해서 각 입력 문자열을 정수 인덱스로 치환해 관리하는 식으로 구현했다.

2. 집합에 속하는 원소의 개수를 구해야 함
   유니온 파인드를 수행하고 끝나는 게 아니라, 방금 유니온한 집합에 속해 있는 원소의 개수를 출력해야 했다. 개수 구하는 건 한 번도 해본 적이 없어서 처음에는 그냥 친구 관계가 생길 때마다 모든 배열 원소를 순회하며 parent 값이 동일한 원소의 개수를 세었는데, 이렇게 하니까 시간 초과가 났다.
   찾아보니까 매번 셀 필요가 없고, 원소의 개수를 저장하는 배열을 따로 생성한 다음 유니온을 수행하는 시점에 합칠 두 집합의 원소 개수 값을 한 쪽으로 합쳐주면 되더라. 이렇게 개수 구하는 문제는 앞으로도 자주 보게 될 것 같아서 기억해둬야겠다고 생각했다.

구현

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

class Main {
    static int[] parent;
    static int[] size;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();
        int T = Integer.parseInt(br.readLine());
        for (int t = 0; t < T; t++) {
            Map<String, Integer> map = new HashMap<>();
            int F = Integer.parseInt(br.readLine());
            parent = new int[F << 1];
            size = new int[F << 1];
            for (int i = 0; i < (F << 1); i++){
                parent[i] = i;
                size[i] = 1;
            }
            
            int index = 0;
            for (int i = 0; i < F; i++) {
                StringTokenizer st = new StringTokenizer(br.readLine());
                String a = st.nextToken();
                String b = st.nextToken();
                if (!map.containsKey(a)) {
                    map.put(a, index++);
                }
                if (!map.containsKey(b)) {
                    map.put(b, index++);
                }
                int parentA = find(map.get(a));
                int parentB = find(map.get(b));
                sb.append(union(parentA, parentB)).append('\n');
            }
        }
        System.out.print(sb.toString());
    }

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

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

        if (a == b) {
            return size[a];
        }

        if (a > b) {
            parent[b] = a;
            size[a] += size[b];
            return size[a];
        }
        
        parent[a] = b;
        size[b] += size[a];
        return size[b];
    }
}