[Java] 백준 4195번: 친구 네트워크

[문제 바로 가기] - 친구 네트워크

유형 : Union-Find

이번주 알고리즘 스터디는 Union-Find 익히기!
마침 최소 스패닝 트리 문제를 여러번 다시 풀었는데도 머리에 들어오지 않았는데 잘된 것 같다. 유파 부시기!🔥

먼저 문제를 이해하는 데 좀 걸렸던 이유가 친구 관계가 생길 때마다, 두 사람의 친구 네트워크에 몇 명이 있는지 구하는 프로그램을 작성하시오. 부분이 이해가 잘 안됐다.

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

예제가 위와 같을 때 출력은 다음과 같다.

2
3
4
2
2
4

이는 첫번째 줄은 Fred와 Barney의 친구 네트워크에 몇명이 존재하는지 출력한 것이다. 즉, 친구 네트워크는 방금 생긴 Fred와 Barney 밖에 없으므로 2명이 출력된다.

두번째 줄에서는 Barney와 Betty가 친구 관계를 맺었으므로 Fred, Barney에 Betty가 추가된 것이다. 친구 네트워크란 친구 관계만으로 이동할 수 있는 사이를 말한다.에 따르면 Barney와 Betty의 친구 네트워크에 있는 사람은 3명이 된다.

💡 아이디어

위에서 이해한 대로 문제를 풀이하려면, Union-Find를 사용할 수 있다.

보통의 Union-Find 풀이는 인덱스를 가지고 parent 배열에 넣는데 이 문제는 String으로 이름이 주어지기 때문에 parent 배열 말고도 HashMap을 사용해주었다.
HashMap에는 이름을 키, index(문제에서는 주어진 순서)를 값으로 넣어주었다.

또한 각 입력마다 친구 네트워크에 있는 친구의 수를 구해줘야 하므로 친구 수를 저장할 배열 size를 선언해주었다. union할 때마다 부모에 자식의 size를 더해주었다.

그리고 여담..

아직 Union Find가 익숙하지 않은지 구현 과정에서 코드가 헷갈렸다.
이번주 꼭 마스터 해야지!

---

풀이

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.Map;
import java.util.StringTokenizer;

/**
 * 백준 4195번 친구 네트워크
 * - Union Find
 */

public class Main {
	public static int[] parent, size;
	
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st;
		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 * 2];
			size = new int[F * 2];
			
			for (int i = 0; i < parent.length; i++) {
				parent[i] = i;
				size[i] = 1;
			}
			
			int index = 0;
			
			for (int i = 0; i < F; i++) {
				st = new StringTokenizer(br.readLine());
				String f1 = st.nextToken();
				String f2 = st.nextToken();
				
				if (!map.containsKey(f1)) {
					map.put(f1, index++);
				}
				
				if (!map.containsKey(f2)) {
					map.put(f2, index++);
				}
				
				int i1 = map.get(f1);
				int i2 = map.get(f2);
				
				sb.append(union(i1, i2) + "\n");
			}
		}
		
		System.out.println(sb.toString());
	}
	
	public static int find(int f) {
		if (parent[f] == f) return f;
		return parent[f] = find(parent[f]);
	}
	
	public static int union(int f1, int f2) {
		f1 = find(f1);
		f2 = find(f2);
		
		if (f1 != f2) {
			parent[f2] = f1;
			size[f1] += size[f2];
			return size[f1];
		}
		
		return size[f1];
	}
}