[백준/자바] 4195번: 친구 네트워크

문제

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

풀이

• 어떤 사이트의 친구 관계가 생긴 순서대로 주어 졌을 때, 두 사람의 친구 네트워크에 몇 명이 있는지 출력
• 친구 네트워크란 친구 관계만으로 이동할 수 있는 사이를 말함

유니온 파인드(Union-Find)를 활용하는 문제

저는 처음에 친구 이름이 들어오면, 값이 존재하는지 비교하여 "리스트"에 값을 넣고 해당 인덱스를 가지고 유니온 파인드를 진행했었습니다.
그러나 조회를 할 때마다 O(N)의 시간이 발생해서 그런지? 시간 초과가 발생하더군요..

그래서 Key값을 바로 확인할 수 있는 Map을 사용했습니다.

			for (int i = 0; i < f; i++) {
				StringTokenizer st = new StringTokenizer(br.readLine());
				String a = st.nextToken();
				String b = st.nextToken();
				
				if (!friends.containsKey(a)) {
					friends.put(a, ++idx);
				}
				
				if (!friends.containsKey(b)) {
					friends.put(b, ++idx);
				}
				
				int size = union(friends.get(a), friends.get(b));
				sb.append(size).append("\n");
			}

• a와 b의 값이 존재하지 않으면 Map에 넣어준 후 Union 연산을 진행했습니다.

	private static void makeSet() {
		p = new int[f * 2 + 1];
		s = new int[f * 2 + 1];
		for (int i = 1; i <= f * 2; i++) {
			p[i] = i;
			s[i] = 1;
		}
	}

• 그러기 위해서 부모, 크기 배열을 선언해 주었습니다.
• 한 번 입력에 값이 2개씩 들어오므로 최악의 경우 2 F개가 발생할 것을 대비해 배열의 크기를 `f 2 + 1`로 초기화했습니다.

	private static int find(int x) {
		if (p[x] == x) return x; // 부모가 자기 자신이면 자신 리턴
		return p[x] = find(p[x]);
	}

• 부모를 찾아 주는 Find 메서드입니다.

	private static int union(int a, int b) {
		int ra = find(a), rb = find(b); // a의 부모, b의 부모
		
		if (ra == rb) return s[ra]; // 부모가 같다면 바로 리턴
		
        // rb를 ra의 밑으로 넣을 것이기 때문에
        // 크기가 역전된 상황이라면 swap
		if (s[ra] < s[rb]) {
			int tmp = ra;
			ra = rb;
			rb = tmp;
		}
		
        // rb의 부모를 ra로 설정(합집합)
		p[rb] = ra;
		s[ra] += s[rb]; // ra 밑에 rb가 들어왔으므로 ra의 크기에 rb 크기만큼 증가
		return s[ra];
	}

• 부모를 합치는 Union 연산입니다.

코드

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

public class Main_4195 {
	static StringBuilder sb = new StringBuilder();
	static int f;
	static int[] p, s;
	static Map<String, Integer> friends;
	
	private static void makeSet() {
		p = new int[f * 2 + 1];
		s = new int[f * 2 + 1];
		for (int i = 1; i <= f * 2; i++) {
			p[i] = i;
			s[i] = 1;
		}
	}
	
	private static int find(int x) {
		if (p[x] == x) return x;
		return p[x] = find(p[x]);
	}
	
	private static int union(int a, int b) {
		int ra = find(a), rb = find(b);
		
		if (ra == rb) return s[ra];
		
		if (s[ra] < s[rb]) {
			int tmp = ra;
			ra = rb;
			rb = tmp;
		}
		
		p[rb] = ra;
		s[ra] += s[rb];
		return s[ra];
	}
	
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int t = Integer.parseInt(br.readLine());
		for (int tc = 1; tc <= t; tc++) {
			f = Integer.parseInt(br.readLine());
			
			friends = new HashMap<>();
			makeSet();
			int idx = 0;
			
			for (int i = 0; i < f; i++) {
				StringTokenizer st = new StringTokenizer(br.readLine());
				String a = st.nextToken();
				String b = st.nextToken();
				
				if (!friends.containsKey(a)) {
					friends.put(a, ++idx);
				}
				
				if (!friends.containsKey(b)) {
					friends.put(b, ++idx);
				}
				
				int size = union(friends.get(a), friends.get(b));
				sb.append(size).append("\n");
			}
		}
		System.out.println(sb.toString());
	}

}