[BOJ 4195] 친구 네트워크 (Java)

문제

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

4195번: 친구 네트워크

민혁이는 소셜 네트워크 사이트에서 친구를 만드는 것을 좋아하는 친구이다. 우표를 모으는 취미가 있듯이, 민혁이는 소셜 네트워크 사이트에서 친구를 모으는 것이 취미이다. 어떤 사이트의 친구 관계가 생긴 순서대로 주어졌을 때, 두 사람의 친구 네트워크에 몇 명이 있는지 구하는 프로그램을 작성하시오. 친구 네트워크란 친구 관계만으로 이동할 수 있는 사이를 말한다.

풀이

union-find 기법을 사용하여 풀 수 있다.

 

1. parent 정보와 네트워크 size 정보 담을 자료 구조를 생성한다.

        parent = new HashMap<String, String>(); // 루트 노드 정보
        setSize = new HashMap<String, Integer>(); // 집합의 사이즈

 

2. 입력 전에는 어떤 값이 들어오는지 알 수 없으므로, 입력 받을 때마다 초기값을 할당한다. 

            // 입력 받을 때 초기화
            if(!setSize.containsKey(a)) setSize.put(a, 1);
            if(!setSize.containsKey(b)) setSize.put(b, 1);
            if(!parent.containsKey(a)) parent.put(a, a);
            if(!parent.containsKey(b)) parent.put(b, b);

 

3. 만약 두 친구가 동일한 네트워크에 있지 않을 경우, 동일한 네트워크로 합친다. 

            if(find(a) != find(b)){
                union(a, b);
            }

 

find함수 내 path-compression기법을 통해 각 노드의 parent에 루트 노드 값을 설정하고, 루트 노드 값을 얻어온다. a, b 중 속한 집합의 사이즈가 큰 쪽으로 Union 한다. 각 네트워크의 루트 노드의 집합 사이즈를 갱신 한다.

    private static String find(String a){
        if(parent.get(a) == a) return a;
        parent.put(a, find(parent.get(a)));
        return parent.get(a);
    }
    
    private static void union(String a, String b){
        a = find(a);
        b = find(b);

        if(setSize.get(a) >= setSize.get(b)){
            parent.put(b, a);
            setSize.put(a, setSize.get(a) + setSize.get(b));
            setSize.put(b, 0);
        } else {
            parent.put(a, b);
            setSize.put(b, setSize.get(b) + setSize.get(a));
            setSize.put(a, 0);
        }
    }

 

4. 네트워크의 사이즈를 출력한다. 

// a, b 친구 관계를 맺었으므로 한 명의 루트 노드에 대한 사이즈만 출력하면 됨
bw.write(setSize.get(find(a)) + "\n");

예제

2
3
Fred Barney
Barney Betty
Betty Wilma
3
Fred Barney
Betty Wilma
Barney Betty

 

첫 번째 Test Case

1. Fred Barney 입력

<Parent>

Key	Value
Fred	Fred
Barney	Fred

<setSize>

Key	Value
Fred	2
Barney	0

<Tree>

 

2. Barney Betty 입력

<Parent>

Key	Value
Fred	Fred
Barney	Fred
Betty	Fred

<setSize>

Key	Value
Fred	3
Barney	0
Betty	0

<Tree>

3. Betty Wilma 입력

<Parent>

Key	Value
Fred	Fred
Barney	Fred
Betty	Fred
Wilma	Fred

<setSize>

Key	Value
Fred	4
Barney	0
Betty	0
Wilma	0

<Tree>

Answer:
2
3
4

 

두 번째 Test Case

1. Fred Barney 입력

<Parent>

Key	Value
Fred	Fred
Barney	Fred

<setSize>

Key	Value
Fred	2
Barney	0

<Tree>

 

2. Betty Wilma 입력

<Parent>

Key	Value
Fred	Fred
Barney	Fred
Betty	Betty
Wilma	Betty

<setSize>

Key	Value
Fred	2
Barney	0
Betty	2
Wilma	0

<Tree>

3. Barney Betty 입력

<Parent>

Key	Value
Fred	Fred
Barney	Fred
Betty	Fred
Wilma	Fred

<setSize>

Key	Value
Fred	4
Barney	0
Betty	0
Wilma	0

<Tree>

Answer:
2
2
4

소스코드

package unionfind;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.HashMap;
import java.util.StringTokenizer;

public class BOJ4195 {

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        int T = Integer.parseInt(br.readLine());
        for(int i = 0; i < T; i++){
            solution(br, bw);
        }
        bw.flush();
    }

    public static HashMap<String, String> parent;
    public static HashMap<String, Integer> setSize;
    private static void solution(BufferedReader br, BufferedWriter bw) throws Exception {
        int N = Integer.parseInt(br.readLine());
        parent = new HashMap<String, String>(); // 루트 노드 정보
        setSize = new HashMap<String, Integer>(); // 집합의 사이즈

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

            // 입력 받을 때 초기화
            if(!setSize.containsKey(a)) setSize.put(a, 1);
            if(!setSize.containsKey(b)) setSize.put(b, 1);
            if(!parent.containsKey(a)) parent.put(a, a);
            if(!parent.containsKey(b)) parent.put(b, b);

            if(find(a) != find(b)){
                union(a, b);
            }
            // a, b 친구 관계를 맺었으므로 한 명의 루트 노드에 대한 사이즈만 출력하면 됨
            bw.write(setSize.get(find(a)) + "\n");
        }
    }

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

        if(setSize.get(a) >= setSize.get(b)){
            parent.put(b, a);
            setSize.put(a, setSize.get(a) + setSize.get(b));
            setSize.put(b, 0);
        } else {
            parent.put(a, b);
            setSize.put(b, setSize.get(b) + setSize.get(a));
            setSize.put(a, 0);
        }
    }

    private static String find(String a){
        if(parent.get(a) == a) return a;
        parent.put(a, find(parent.get(a)));
        return parent.get(a);
    }

}