[백준 / BOJ] 19541번 : 역학 조사 - JAVA

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

19541번: 역학 조사

기본 아이디어

문제를 보고 가장 먼저 떠올린 것을 가장 마지막 모임부터 거꾸로 올라가면서 병이 전염된 모임을 찾는 것이었다. 모든 참가자가 병에 걸려있다면 해당 모임에서 전염병에 걸렸다고 가정하고, 한 명이라도 병에 걸려 있지 않으면 그 모임이 시작하기 전까지는 해당 모임의 참가자가 병에 걸리지 않은 것이다.

 

즉, 정리하면 다음과 같다.

1. N번째 모임의 참가자가 모두 전염병에 걸려있다면 해당 모임이 시작하기 전 해당 참가자들은 모두 전염병에 걸려 있었다고 가정한다.
2. N번째 모임의 참가자 중 한 명이라도 전염병에 걸려 있지 않다면 해당 모임이 시작하기 전 해당 참가자들은 모두 전염병에 걸려있지 않았다고 가정한다.

이를 통해 모임을 M - 1번째부터 0번째까지 거슬러 올라가며 최초 감염자를 찾고, 만일 최초 감염자가 없다면 NO를 출력했다. 단, 이렇게만 코드를 짜면 시작하자마자 틀리고 만다.

 

반례

내가 겪은 반례는 다음과 같다.

#1 
3 1
2 1 2
1 0 1

정답 :
NO

#2
3 1
2 1 2
1 1 1

정답 :
YES
1 1 1

#3
4 2
2 1 2
2 3 4
1 0 1 1

정답 :
NO

#4
4 1
3 1 2 3
0 0 0 0

정답 :
YES
0 0 0 0

 

위 반례를 고려하여 다음과 같은 절차를 추가했다.

3. N번째 모임의 참여자가 모두 전염병에 걸렸다면, 참여자가 전염병에 걸린 모임을 나타내는 order 배열의 값을 N으로 업데이트 한다. (order 배열은 최초에 -1로 초기화한다.)
    3 - 1. 최종 감염자인 사람 i의 order[i] 값이 -1이고, 모임에 참여한 적이 있다면(v[i]가 true라면) 해당 참여자가 전염병에 걸린 모임을 찾을 수 없으므로 "NO"를 출력한다.

3 - 1 조건을 넣은 이유는, 최종 감염자이지만 모임에 한 번도 참여를 안했다면 그냥 처음부터 걸려 있었을 것이기 때문이다. 따라서 모임에 참여를 한 번도 하지 않은 사람은 최초 감염자를 찾을 때 제외하고 생각해야 한다.

 

전체 코드

public class Main {

	public static void main(String[] args) throws IOException {

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		
		int N = Integer.parseInt(st.nextToken());
		int M = Integer.parseInt(st.nextToken());
		
		List<Integer>[] list = new ArrayList[M];
		boolean[] v = new boolean[N];
		for (int i = 0; i < M; i++) {
			list[i] = new ArrayList<>();
			st = new StringTokenizer(br.readLine());
			int K = Integer.parseInt(st.nextToken());
			for (int j = 0; j < K; j++) {
				int a = Integer.parseInt(st.nextToken()) - 1;
				list[i].add(a);
				v[a] = true;
			}
		}
		
		int[] covid = new int[N];
		int[] origin = new int[N];
		st = new StringTokenizer(br.readLine());
		for (int i = 0; i < N; i++) {
			covid[i] = Integer.parseInt(st.nextToken());
			origin[i] = covid[i];
		}
		
		int[] order = new int[N];
		Arrays.fill(order, -1);
		for (int i = M - 1; i >= 0; i--) {
			boolean flag = true;
			for (int j : list[i]) {
				if(origin[j] == 0) {
					for (int k : list[i]) {
						origin[k] = 0;
					}
					flag = false;
					break;
				}
			}
			if(flag) {
				for (int j : list[i]) {
					order[j] = i;
				}
			}
		}
		
		StringBuilder sb = new StringBuilder();
		sb.append("YES\n");
		for (int i = 0; i < N; i++) {
			if(covid[i] == 1 && order[i] == -1 && v[i]) {
				System.out.println("NO");
				return;
			}
			sb.append(origin[i]).append(" ");
		}
		System.out.println(sb);
		
	}

}