[백준] 25187번 : 고인물이 싫어요 (C++)

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

문제

재형이는 청정수를 좋아하고 고인물을 싫어한다. 오늘도 청정수를 구하기 위해 물탱크들이 있는 마을에 방문한다.

마을에는 N개의 물탱크가 존재하고, 각 물탱크는 청정수 또는 고인물을 저장하고 있다. 그리고 물탱크는 공급의 편의를 위해 M개의 파이프로 서로 연결되어 있다.

청정수를 얻기 위해 K번 물탱크에 방문했을 때, K번 물탱크와 K번 물탱크에서 0개 이상의 파이프를 거쳐 이동 가능한 물탱크 중, 청정수가 담긴 물탱크의 수가 고인물이 담긴 물탱크의 수보다 더 많은 경우 청정수를 얻을 수 있다.

방문할 예정인 물탱크에 대한 정보가 주어질 때마다, 청정수를 얻을 수 있는지 구해보자. 

입력

첫째 줄에 물탱크의 수 N(1≤N≤100000)과 파이프의 수 M(0≤M≤200000), 그리고 물탱크에 방문할 횟수 Q(1≤Q≤100000)가 주어진다. 

둘째 줄에 1번부터 N번 물탱크까지 순서대로 들어있는 물의 종류가 주어진다. 청정수는 1, 고인물은 0으로 주어진다. 

다음 3번째부터 M+2번째 줄까지 파이프가 연결하는 두 물탱크의 번호 u,v(1≤u,v≤N,u≠v)가 주어진다. 같은 두 물탱크를 연결하는 파이프가 여러 개 존재할 수 있다. 

 M+3번째부터 M+Q+2번째 줄까지 방문할 물탱크의 번호 K(1≤K≤N)가 주어진다.  

출력

 Q개의 줄에 각 물탱크에 방문했을 때 청정수를 얻을 수 있다면 1을, 아니면 0을 주어진 정보 순서대로 출력한다. 

예제 입력 1 

5 3 3
1 0 1 1 0
1 2
3 4
4 5
1
5
4

예제 출력 1 

0
1
1

첫번째 쿼리는 1번 물탱크와 연결된 물탱크는 2번 물탱크 하나이므로, 고인물이 담긴 물탱크와 청정수가 담긴 물탱크의 수가 같다. 

두번째 쿼리는 5번 물탱크와 연결된 물탱크는 3, 4번 물탱크이므로, 고인물이 담긴 물탱크보다 청정수가 담긴 물탱크가 더 많다. 

예제 입력 2 

5 5 3
1 0 1 1 0
1 2
1 3
1 4
2 3
3 4
1
5
4

예제 출력 2 

1
0
1

---

트리 순회문제이다. 그래서 처음에 DFS로 구현했다. 그런데 시간제한이 1초라 TLE에 걸려벌임.

 

이 문제는 Union Find를 이용해 풀이해야한다.

Union Find

// 초기화

void Init(int n){

    for(int i = 1 ; i <= n ; i++){

        P[i]=i;

    }

}

 

// RootNode 찾기

int Find(int v){

    if(v==P[v]) return v;

    return Find(P[v]);

}

 

// 두 노드를 같은 집합으로 합치기

void Union(int u, int v){

    u = Find(u); v = Find(v);

    if(u != v) P[u] = v;

}

 

각 노드에 대해 Parent 배열을 구성해 Root를 찾아간다고 생각하면 된다.

 

문제는 다음과 같이 구성되었을 때 Find를 하는데 O(h)나 걸린다는것.

 

 

 

 

 

 

 

 

이렇게 트리가 구성되면 O(1)만에 Root를 찾을 수 있다.

 

 

 

 

 

 

 

그래서 트리를 잘압축하는게 중요한데, Union Find 트리는 밑의 세가지 중 하나의 방식을 채택하면 O(MlogN)안에 탐색이 가능함이 보장된다.

1.  Union by Rank
   높이가 낮은 트리를 높은 트리 아래로 넣는 방법
   
   int P[100000]; int R[1000000];

   void Init(int n){

       for(int i = 1 ; i <= n ; i++){

           P[i]=i; R[i] = 1;

       }

   }

   int Find(int v){

       if(v==P[v]) return v;

       return P[v] = Find(P[v]);

   }

   void Union(int u, int v){

       u = Find(u); v = Find(v);

       if(u==v) return;

       if(R[v] < R[u]) swap(u,v);

       P[u] = v;

       if(R[u] ==R[v]) P[v]++;     //높이가 같을 경우 트리높이가 1 높아진다.

   }
   
   R배열을 추가로 만들어 높이를 비교한 뒤 높이가 낮은 트리를 높은 트리 안으로 넣는다.
   
   
2. Union by Size
   정점이 적은 트리를 많은 트리 아래로 넣는다.
   
   int P[100000]; int N[100000];

   void Init(int n){

       for(int i = 1 ; i <= n ; i++){

           P[i]=i;

       }

   }

   int Find(int v){

       if(v==P[v]) return v;

       return Find(P[v]);

   }

   void Union(int u, int v){

       u = Find(u); v = Find(v);

       if(u==v) return;

       if(N[v] < N[u]) swap(u,v);

       P[u] = v;

       N[v] += N[u];

   }
3. 경로 압축
   
   void Init(int n){

       for(int i = 1 ; i <= n ; i++){

           P[i]=i;

       }

   }

   int Find(int v){

       if(v==P[v]) return v;

       return P[v] = Find(P[v]);

   }

   void Union(int u, int v){

       u = Find(u); v = Find(v);

       if(u != v) P[u] = v;

   }
   
   Find로 루트를 찾을 때 루트까지의 경로상에 있는 노드를 모두 루트 밑의 자식으로 붙여버림

---

Union Find로 Disjoint Set을 구한뒤에 집합별로 고인물인지 청정수인지 계산해서 출력

#include <bits/stdc++.h>

using namespace std;

int n, m, q, num;

int V[101010];

int res[101010];

int P[101010];

void Init(int n){

    for(int i = 1 ; i <= n ; i++){

        P[i]=i;

    }

}

int Find(int v){

    if(v==P[v]) return v;

    return P[v] = Find(P[v]);

}

void Union(int u, int v){

    u = Find(u); v = Find(v);

    if(u != v) P[u] = v;

}

void sol(){

    for(int s = 1 ; s <= n ; s++){

        if(V[s]==1){

            res[Find(s)]++;

        }

        else{

            res[Find(s)]--;

        }

    }

}

int main(void){

    ios_base::sync_with_stdio(0); cin.tie(0);

    cin>>n>>m>>q;

    Init(n);

    for(int i = 1 ; i <= n ; i++){

        cin>>V[i];

    }

    for(int i = 0 ; i < m ; i++){

        int s, e;

        cin>>s>>e;

        Union(s,e);

    }

    sol();

    for(int i = 0 ; i < q ; i++){

        int s; cin>>s;

        if(res[P[s]]>0) cout<<1<<"\n";

        else cout<<0<<"\n";

    }

}