백준 4195 친구 네트워크 Java

알고리즘 분류 : 자료 구조, 해시를 사용한 집합과 맵, 분리 집합

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

요약

기존의 유니온 파인드 알고리즘은 입력이 정수로 주어지는데 이번 문제는 문자열이 주어졌다. 평소 Map을 사용한 경험이 있다면 문자열을 어렵지 않게 정수로 변환할 수 있다. 문제에서 전체 사람 수가 주어지지 않아 배열 대신 List를 활용할 수 있는데 이 경우 find의 경로 최적화에서 List의 set 함수를 주의해야 한다.

풀이

친구 네트워크는 친구 관계만으로 이동할 수 있는 사이라고 한다. 풀어서 말하면 A의 친구의 친구의 친구…가 B라면 A와 B는 동일한 친구 네트워크이다. 따라서 그래프 내 연결된 그룹의 원소 수를 구하는 문제이므로 유니온 파인드를 통한 구현을 생각할 수 있다.

다만 이 문제에서 기존의 유니온 파인드 알고리즘에서 추가적으로 해야 할 작업은 다음 두가지이다.

1. 이름을 int 인덱스로 변환할 수 있는 기능
2. 친구 관계가 맺어질 때 친구 네트워크의 사람 수 출력
3. 전체 사람 수를 알 수 없으므로 List로 구현

1번은 Map<String, Integer>로 어렵지 않게 구현할 수 있다. 2번은 병합 연산 시 두 그룹의 원소 수를 합한 후 출력하는 방식으로 구현할 수 있다. 3번은 전체 친구 관계 수가 10만이므로 친구 관계에서 모두 다른 사람이 입력으로 들어올 경우 최대 20만명까지 들어올 수 있다. 따라서 크기가 20만인 배열로도 해결 가능하지만 불필요한 메모리 공간을 줄이기 위해 ArrayList를 사용하였다.

코드

public class Main {
  static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
  public static void main(String[] args) throws Exception {
    int cases = Integer.parseInt(br.readLine());

    for (int i = 0; i < cases; i++) {
      oneCase();
    }
  }

  private static void oneCase() throws Exception {
    int m = Integer.parseInt(br.readLine());
    UnionFind uf = new UnionFind();

    for (int i = 0; i < m; i++) {
      String[] line = br.readLine().split(" ");
      uf.merge(line[0], line[1]);
    }

    uf.printAll();
  }

  private static class UnionFind {
    private List<Integer> parent = new ArrayList<>();
    private List<Integer> rank = new ArrayList<>();
    private List<Integer> rootNodeGroupNumber = new ArrayList<>();

    private Map<String, Integer> nameIndexInfo = new HashMap<>();

    private int userCount = 0;

    private StringBuilder output = new StringBuilder();

    private int find(String name) {
        //맵에 정보가 없으면 새로운 사람이므로 추가
      if (!nameIndexInfo.containsKey(name)) {
        nameIndexInfo.put(name, userCount);
        parent.add(userCount);
        rank.add(1);
        rootNodeGroupNumber.add(1);
        userCount++;
      }
    
      //사람 이름에서 인덱스 가져오기
      int index = nameIndexInfo.get(name);
      return find(index);
    }

    private int find(int index) {
        //루트 노드에 도달하면
      if (index == parent.get(index)) return index;
      
      //경로 압축
      parent.set(index, find(parent.get(index)));
      return parent.get(index);
    }

    public void merge(String a, String b) {
      int findA = find(a);    int findB = find(b);
      
      //이미 서로 동일한 친구 네트워크이면 병합하지 않기
      if (findA == findB) {
        output.append(rootNodeGroupNumber.get(findA)).append('\n');
        return;
      }

      //A의 높이가 더 높으면 A와 B를 교체하여 항상 더 작은 트리가 큰 트리에 종속되도록 함
      if (rank.get(findA) > rank.get(findB)) {
        int temp = findA;
        findA = findB;
        findB = temp;
      }
      
      //트리 하나를 종속시킨다.
      parent.set(findA, findB);
      
      //두 트리(네트워크)의 사람 수를 더하고 갱신한다.
      int totalGroupNumber = rootNodeGroupNumber.get(findA) + rootNodeGroupNumber.get(findB);
      output.append(totalGroupNumber).append('\n');
      rootNodeGroupNumber.set(findA, totalGroupNumber);
      rootNodeGroupNumber.set(findB, totalGroupNumber);
      
      //두 트리의 높이가 같으면 합한 트리의 높이는 1 증가한다.
      if (rank.get(findA).equals(rank.get(findB))) rank.set(findB, rank.get(findB) + 1);
    }

    public void printAll() {
      System.out.print(output);
    }
  }
}

기존 토대는 유니온 파인드와 동일하므로 이번 문제에 필요한 추가적인 코드는 다음과 같다.

1. 이름을 int 인덱스로 변환

    private int find(String name) {
      if (!nameIndexInfo.containsKey(name)) {
          ...
      }

      int index = nameIndexInfo.get(name);
      return find(index);
    }

find 연산을 문자열로 받는다. 이 문자열을 Map에서 인덱스로 얻은 후 기존의 인덱스를 활용한 find 연산으로 넘긴다.

1. 병합 연산 시 두 그룹의 수를 합한 후 출력하기

  public void merge(String a, String b) {
    ...
    int totalGroupNumber = rootNodeGroupNumber.get(findA) + rootNodeGroupNumber.get(findB);
    output.append(totalGroupNumber).append('\n');
    rootNodeGroupNumber.set(findA, totalGroupNumber);
    rootNodeGroupNumber.set(findB, totalGroupNumber);
  }

a와 b의 그룹이 같으면 바로 a나 b 그룹의 사람 수를 반환하면 된다. a와 b의 그룹이 다른 경우 a와 b의 그룹 수를 더하고 갱신하면 된다.

1. 전체 사람 수를 알 수 없으므로 List로 구현

  //아래 두 코드는 다른 결과가 나온다!!!!!!
  return parent[index] = find(parent[index]);

  return parent.set(index, find(parent.get(index)));

대부분 코드는 쉽게 List로 변환할 수 있지만 주의해야 할 점이 있다. 경로 압축 시 위의 코드를 아래로 변경했다간 틀렸습니다 또는 메모리 초과라는 결과 창을 볼 것이다. List의 set은 변경 이전의 값을 반환하기 때문에 find 함수의 결과로 parent.get(index)를 반환하게 된다. 그러므로 다음과 같이 코드를 작성해야 한다.

parent.set(index, find(parent.get(index)));

return parent.get(index);

한 줄 차이로 전혀 다른 결과를 가져올 수 있는 것이다.

결론

유니온 파인드를 사용하는 문제임을 알면 크게 어려운 부분은 없다. 다만 위와 같이 기존의 알고리즘에서 여러 요구 사항에 맞춰 수정하다보면 예상하지 못한 실수가 일어날 수 있으므로 주의해야 한다. 특히 유니온 파인드 시 List로 구현할 때 경로 압축에서 set 사용 시 주의할 점은 잊지 말고 기억하자.

Reference

https://www.acmicpc.net/board/view/115376