백준 1953번 팀배분

문제

백준 1953번 팀배분 링크

처음에 문제를 보고 유니온 파인드가 떠올랐다. 그런데 잘보니 유니온이 아니라 다른 팀이 되야하는 연산을 해야하므로 유니온 파인드로 풀기는 어려웠다.

이후 아래 코드를 구현했다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.*;

public class Main {
    static int N;
    static int [] parent;
    public static void main(String[] args) throws Exception{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;
        StringBuilder sb = new StringBuilder();

        st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());

        parent = new int[N+1];
        Arrays.fill(parent, 1); // 1과 -1로 구분

        for(int i=1; i<=N; i++){
            st = new StringTokenizer(br.readLine());
            int total = Integer.parseInt(st.nextToken());
            int nowTeam = parent[i];
            for(int k=0; k<total; k++){
                int dislike = Integer.parseInt(st.nextToken());
                parent[dislike] = - nowTeam;
            }
        }

        List<Integer> firstTeam = new ArrayList<>();
        List<Integer> secondTeam = new ArrayList<>();
        for(int i=1; i<=N; i++){
            if(parent[i] == 1) firstTeam.add(i);
            else secondTeam.add(i);
        }

        Collections.sort(firstTeam);
        Collections.sort(secondTeam);

        sb.append(firstTeam.size()).append("\n");
        for(int num : firstTeam) sb.append(num).append(" ");
        sb.append("\n");

        sb.append(secondTeam.size()).append("\n");
        for(int num : secondTeam) sb.append(num).append(" ");

        System.out.print(sb);
    }
}

그런데 이 코드는 통과할 수 없다. 반례를 구하기 어려워 내가 직접 반례를 찾아야했다.

4
2 2 3
1 1
1 1
1 1

위와 같은 입력을 하면 내가 작성한 코드에서는 아래와 같이 출력된다.

1
4 
3
1 2 3 

1번은 분명 2, 3번과 같은팀이 되면 안되는데 말이다,,,,

이유는 간단하다. 내가 1번을 보고 2번과 3번을 -1로 초기화 했을 떄 4번은 1인 상태이다. 그런데 4번이 1을 싫어한다는 것을 보는 순간 4번의 값(1)에 -를 붙임 -1을 1의 값으로 설정하게 되어 배열이 -1 -1 -1 1 이 되어버리는 것이다.

따라서 우리는 DFS를 통해 풀어야한다. 왜냐하면 하나가 결정되면 cascade로 결정을 해주고, 또 결정된 것을 cascade로 싫어하는 사람을 -한 값으로 저장해야하기 때문이다. 이러한 과정에서 중요한 것은 visited를 사용해야한다는 것이다. 이미 팀을 배정받은 사람이 다시 팀을 배정받으면 모든 규칙이 엉망이 되어버리기 때문이다.

풀이 방법

1. 싫어하는 사람들을 인접 리스트 배열로 양방향 저장하기

싫어하는 사람들을 인접 리스트 배열로 저장해야한다. 이때 양방향으로 저장해야함에 주의하자.

        for(int i=1; i<=N; i++){
            st = new StringTokenizer(br.readLine());
            int total = Integer.parseInt(st.nextToken());
            for(int k=0; k<total; k++){
                int dislike = Integer.parseInt(st.nextToken());
                arr[i].add(dislike);
                arr[dislike].add(i);
            }
        }

2. 1번은 1번팀에 고정하고, DFS 돌리기

어차피 2개의 팀밖에 주어지지 않으므로 1번은 1번팀에 고정한다.
이후 DFS를 돌린다.

team[1] = 1; // 1번은 1번 팀
        for(int i=1; i<=N; i++){
            if(!visited[i]) DFS(i, team[i]);
        }

static void DFS(int num, int nowTeam){
        visited[num] = true;
        for(int n : arr[num]){
            if(!visited[n]) {
                team[n] = - nowTeam;
                DFS(n, - nowTeam);
            }
        }
    }

3. 순서대로 출력하기

이제 각 팀의 인원을 계산한 뒤 오름차 순으로 출력한다.

int firstTeam = 0; // 첫번째 팀 인원수
        for(int i=1; i<=N; i++){
            if(team[i] == 1) firstTeam++;
        }
        sb.append(firstTeam).append("\n");

        int secondTeam = N - firstTeam; // 두번째 팀 인원수

        // 첫번째 팀 출력
        for(int i=1; i<=N; i++){
            if(team[i] == 1) sb.append(i).append(" ");
        }
        sb.append("\n");

        sb.append(secondTeam).append("\n");
        // 두번째 팀 출력
        for(int i=1; i<=N; i++){
            if(team[i] == -1) sb.append(i).append(" ");
        }

        System.out.print(sb);

정답 코드

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.*;

public class p1953 {
    static int N;
    static int [] team;
    static List<Integer> [] arr;
    static boolean[] visited;

    public static void main(String[] args) throws Exception{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;
        StringBuilder sb = new StringBuilder();

        st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());

        arr = new List[N+1];
        team = new int[N+1];
        visited = new boolean[N+1];
        for(int i=1; i<=N; i++) arr[i] = new ArrayList<>();


        for(int i=1; i<=N; i++){
            st = new StringTokenizer(br.readLine());
            int total = Integer.parseInt(st.nextToken());
            for(int k=0; k<total; k++){
                int dislike = Integer.parseInt(st.nextToken());
                arr[i].add(dislike);
                arr[dislike].add(i);
            }
        }

        team[1] = 1; // 1번은 1번 팀
        for(int i=1; i<=N; i++){
            if(!visited[i]) DFS(i, team[i]);
        }

        int firstTeam = 0; // 첫번째 팀 인원수
        for(int i=1; i<=N; i++){
            if(team[i] == 1) firstTeam++;
        }
        sb.append(firstTeam).append("\n");

        int secondTeam = N - firstTeam; // 두번째 팀 인원수

        // 첫번째 팀 출력
        for(int i=1; i<=N; i++){
            if(team[i] == 1) sb.append(i).append(" ");
        }
        sb.append("\n");

        sb.append(secondTeam).append("\n");
        // 두번째 팀 출력
        for(int i=1; i<=N; i++){
            if(team[i] == -1) sb.append(i).append(" ");
        }

        System.out.print(sb);

    }
    static void DFS(int num, int nowTeam){
        visited[num] = true;
        for(int n : arr[num]){
            if(!visited[n]) {
                team[n] = - nowTeam;
                DFS(n, - nowTeam);
            }
        }
    }
}