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

고범수·2023년 3월 15일
0

Problem Solving

목록 보기
28/31

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

문제 설명

어떤 사이트의 친구 관계가 생긴 순서대로 주어졌을 때, 두 사람의 친구 네트워크에 몇 명이 있는지 구하는 프로그램을 작성하는 문제이다. '친구1 친구2' 이런 관계가 주어지면 친구1의 친구 네트워크에 몇 명, 친구2의 친구 네트워크에 몇 명인지 출력하면 된다. 친구 네트워크란 친구 관계만으로 이동할 수 있는 사이를 말한다.

문제 풀이

친구1이 친구네트워크1에 속하고 친구2가 친구네트워크2에 속할 때, 친구3이 친구1, 친구2와 친구라면 친구3은 어느 친구네트워크에 속하게 되는가? 를 생각하면 이 문제가 서로소 집합(Disjoint Set) 문제임을 알 수 있다.

위 질문의 답은 "친구1 친구2 친구3 모두 한 네트워크에 속한다."가 된다. 다시 말해, 교집합이 생기지 않고 두 집합이 하나의 집합으로 합쳐지기 때문에 어느 원소 하나는 하나의 집합에만 속하는 서로소 집합이 된다.

서로소 집합은 parents[] 로 union-find 연산을 구현하여 생성할 수 있다. 그러나 복잡도를 개선하기 위해서는 rank 개념과 path compression을 수행해야 한다.

친구 관계가 주어질 때마다 각 집합에 속하는 원소에 개수를 출력해야 하므로 각 서로소 집합의 대표자(representative)로 접근가능한 size[] 에 그 집합의 원소의 수를 계산해 놓는다.

마지막으로, 보통 서로소 집합은 배열과 idx로 구현하는데 이 문제는 원소가 String이다. 따라서 Map<String, String>으로 parents[]를 대신하거나 Map<String, Integer>로 String을 idx로 매핑해서 구현하면 된다.

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

public class Main {
    static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
    static StringTokenizer st;

    static int f;
    static final int MAX = 200_000;
    static Map<String, Integer> nameToIdx;
    static DisSet disSet;

    static class DisSet {
        int[] parents;
        int[] ranks;
        int[] size;

        public DisSet() {
            this.parents = new int[MAX];
            this.ranks = new int[MAX];
            this.size = new int[MAX];

            for (int i = 0; i < MAX; i++) {
                parents[i] = i;
                size[i] = 1;
            }
        }

        public int find(int idx) {
            if (parents[idx] == idx) {
                return idx;
            }

            return parents[idx] = find(parents[idx]);
        }

        public boolean union(int a, int b) {
            int setA = find(a);
            int setB = find(b);

            if (setA == setB) {
                return false;
            }

            if (ranks[setA] > ranks[setB]) {
                parents[setB] = setA;
                size[setA] += size[setB];
            } else if (ranks[setA] < ranks[setB]) {
                parents[setA] = setB;
                size[setB] += size[setA];
            } else {
                parents[setB] = setA;
                ranks[setA]++;
                size[setA] += size[setB];
            }

            return true;
        }

        public int getSize(int idx) {
            int set = find(idx);

            return size[set];
        }
    }

    public static void main(String[] args) throws Exception {
        int testNum = Integer.parseInt(br.readLine());
        for (int testCnt = 1; testCnt <= testNum; testCnt++) {
            solution();
        }

        bw.flush();
    }

    static void solution() throws Exception {
        f = Integer.parseInt(br.readLine());

        nameToIdx = new HashMap<>();
        disSet = new DisSet();

        int idx = 0;
        for (int i = 0; i < f; i++) {
            st = new StringTokenizer(br.readLine());

            String name1 = st.nextToken();
            String name2 = st.nextToken();
            if (!nameToIdx.containsKey(name1)) {
                nameToIdx.put(name1, idx++);
            }

            if (!nameToIdx.containsKey(name2)) {
                nameToIdx.put(name2, idx++);
            }

            disSet.union(nameToIdx.get(name1), nameToIdx.get(name2));

            bw.write(Integer.toString(disSet.getSize((nameToIdx.get(name1)))));
            bw.newLine();
        }
    }

}

0개의 댓글