[BOJ/Platinum 5] 백준 19541 역학 조사(C++)

문제 링크

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

19541번: 역학 조사

문제

2020년, 신종 전염병이 유행하여 UCPC국 질병관리본부에서 역학 조사를 하고 있다. UCPC국의 인구는 총 N명이며 각각 1, 2, ... ,N번의 주민번호가 붙어있다.

질병관리본부는 지금까지 M개의 모임이 있었다는 사실을 파악했다. 각 모임은 k명이 참여했고 해당 모임은 주민번호가 a1, a2, ..., ak인 사람들이 참여했다.

전염병은 밀접하고 밀폐된 공간에서만 전염되기 때문에 반드시 모임 안에서만 전염된다. 전염병이 전파되는 규칙은 다음과 같다.

• 모임에 참여한 사람들 중 한 명 이상의 사람이 전염병에 감염되어 있었다면 모임에 참여한 모든 사람들이 전염병에 감염된다.
• 모임에 전염병에 감염된 사람이 없다면 아무 일도 일어나지 않는다.

질병관리본부는 확보한 자료를 가지고 초기 감염자들을 예측하려고 한다. 모임의 정보 및 M개의 모임이 끝나고 나서 전염병에 감염된 사람의 정보가 주어지면 첫 번째 모임을 하기 전에 감염되어 있던 사람을 역추적하는 프로그램을 작성하여라. 위 규칙 이외의 경로로 전염병이 전파되거나 전염병이 치료되는 경우는 없다고 간주한다.

입력

첫 번째 줄에 사람의 수 N, 모임의 수 M (2≤N≤100 000, 1≤M≤100 000)이 주어진다.

두 번째 줄부터 M개의 줄에는 모임의 정보가 시간 순으로 주어진다. 각 줄에는 각 모임에 참여하는 사람의 수 k (2≤k≤N)와 모임에 참여한 사람의 주민번호 ai(1≤ai≤N, ai≠aj)가 주어진다. 여러 모임이 동시에 진행되는 경우는 없다.

마지막 줄에는 N명의 사람들에 대한 감염 정보가 주어진다. 마지막 모임이 끝나고 주민번호가 i인 사람이 전염병에 감염되었다면 1을, 그렇지 않다면 0이 주어진다. 감염된 사람이 없을 수 있음에 유의하여라.

 k들의 합은 1,000,000을 넘지 않는다.

출력

만약 모임을 하기 전에 감염된 사람을 역추적할 수 없다면 NO를 출력한다.

그렇지 않다면 첫 번째 줄에 YES를 출력하고 두 번째 줄에 감염 정보를 의미하는 정수 N개를 공백으로 구분하여 출력한다. 이 중 i번째 수는 주민번호가 i인 사람이 첫 번째 모임이 시작되기 전에 전염병에 감염되어 있었으면 1이며, 아니면 0이다. 가능한 감염 상태가 여러 개이면, 그중 아무 것이나 출력하면 된다.

예제 입력 1

7 3
3 1 2 3
3 3 4 5
3 5 6 7
0 0 1 1 1 1 1

예제 출력 1

YES
0 0 0 1 1 1 1

예제 입력 2

7 3
3 1 2 3
3 3 4 5
3 5 6 7
1 0 1 0 1 0 1

예제 출력 2

NO

 

알고리즘 분류

• 구현
• 그리디 알고리즘

풀이

모임 당 10만명이 소속되어 있는 10만개의 모임을 다 탐색하기 위해서 시간이 충분히 주어졌다.

모임을 우선 역순으로 탐색하여, 한 명이라도 최종 감염자가 없다면 그 모임은 초기 감염자가 존재하지 않는다고 생각할 수 있다. 모임을 모두 탐색하면서 초기 감염자를 분류(예측)하고, 다시 모임을 순서대로 탐색하면서 초기 감염자가 한 명이라도 존재한다면 그 모임 전원을 최종 감염자로 분류(예측)한다.

마지막으로 예측하여 분류한 최종 감염자와 실제로 입력받은 최종 감염자를 비교하여, 한 명이라도 정보가 다르다면 NO, 정보가 일치한다면 YES를 출력하고 실제로 예측한 초기 감염자들을 출력한다.

코드

#include <iostream>
#include <string>
#include <cmath>
#include <set>
#include <unordered_map>
#include <stack>
#include <queue>
#include <vector>
#include <algorithm>
#define FIRST cin.tie(NULL); cout.tie(NULL); ios::sync_with_stdio(false);
#define MAX 100005
#define LL long long
#define INF 1e9

using namespace std;
int N, M;
vector<int> Meeting[MAX];
int visited[MAX]; // 주어진 최종 감염자 데이터
int initial[MAX]; // 초기 감염자 후보
int infected[MAX]; // 최종 감염자
bool answer = true;

int main() {
	FIRST
	cin >> N >> M;
	for (int i = 0; i < M; i++) {
		int K;
		cin >> K;
		for (int j = 0; j < K; j++) {
			int A;
			cin >> A;
			Meeting[i].push_back(A);
		}
	}
	for (int i = 1; i <= N; i++) {
		cin >> visited[i];
		infected[i] = visited[i]; // 최종 감염자의 정보를 저장
		initial[i] = visited[i]; // 초기 감염자 후보를 저장
	}
	// 모임을 역순으로 거슬러 올라간다.
	for (int i = (M - 1); i >= 0; i--) {
		bool Flag = true;
		for (int j = 0; j < Meeting[i].size(); j++) {
			int A = Meeting[i][j];
			if (visited[A] == 0) { // 해당 모임에서 최종 비감염자가 있는 경우
				// 해당 모임의 모든 인원들은 초기 감염자가 될 수 없다.
				Flag = false;
				break;
			}
		}
		if (!Flag) { // 최종 감염자가 한 명이라도 존재하지 않으면 모든 인원들은 초기 감염자가 아님
			for (int j = 0; j < Meeting[i].size(); j++) {
				int A = Meeting[i][j];
				initial[A] = 0;
				visited[A] = 0;
			}
		}
	}
	for (int i = 1; i <= N; i++) {
		visited[i] = initial[i]; // 남은 후보들을 초기감염자라고 가정한다.
	}
	// 모임을 시간 순서대로 거슬러 올라간다.
	for (int i = 0; i < M; i++) {
		bool Flag = false;
		for (int j = 0; j < Meeting[i].size(); j++) {
			int A = Meeting[i][j];
			if (visited[A] == 1) { // 해당 모임에서 초기 감염자 후보가 있는 경우
				Flag = true;
				break;
			}
		}
		if (Flag) { // 초기 감염자 후보가 한 명이라도 존재한다면 모임의 모든 사람이 감염된다(고 가정)
			for (int j = 0; j < Meeting[i].size(); j++) {
				int A = Meeting[i][j];
				visited[A] = 1;
			}
		}
	}
	// 확인한 감염 정보와 최종 감염자의 정보가 일치하지 않는다면 NO 출력
	for (int i = 1; i <= N; i++) {
		if (infected[i] != visited[i]) {
			answer = false;
			break;
		}
	}
	
	if (answer) {
		cout << "YES" << "\n";
		for (int i = 1; i <= N; i++) {
			cout << initial[i] << " ";
		}
	}
	else {
		cout << "NO" << "\n";
	}

	return 0;
}