18267 - Milk Visits

문제

** 오역이 있을 수 있으니 참고 부탁드립니다.

농부 John 은 N (1 ~ 10⁵) 개의 농장을 지을 계획이다. 각 농장은 N-1 개의 길로 연결되어 있고,
농장은 tree 구조를 형성한다. (즉, 모든 농장들은 cycle 을 형성하지 않고, 어떻게든 다른 농장으로 갈 수 있다.)
각 농장은 소를 가지고 있으며, 소의 품종은 Guernsey 또는 Holsein 이다.

Jhon 의 M (1 ~ 10⁵) 명의 친구들은 종종 놀러온다. 친구 i 가 방문하는 동안,
존은 A_i 부터 B_i 까지 가는 유일한(unique) 길을 따라 친구와 함께 걷는다.

그들은 길을 따라 걸으면서 소의 우유도 먹을 수 있다. (john 씨 친절)
그의 친구들은 Guernsey 우유 나 Holstein 우유만 먹을 수 있다.
존의 친구들은 그들이 먹고 싶은 우유를 먹어야만 행복할 수 있다.

방문 후에 어떤 친구들이 행복할 수 있을지 알려주세요!

입력

첫 번째 줄 : N , M
두 번째 줄 : 농장에 있는 소의 종류가 string 으로 ( H or G)
N-1 개 줄 : 연결되어 있는 농장
M 개 줄 : A, B, C (친구 i 가 걸을 시작농장과 끝 농장, 친구 i 의 우유 취향)

5 5
HHGHG
1 2
2 3
2 4
1 5
1 4 H
1 4 G
1 3 G
1 3 H
5 5 H

출력

happy 할 수 있는 친구면 1을 아니면 0을 출력한다.

풀이

생각을 해보자

• 존씨 친구가 걷는 길 도중에 H, G 종류의 소가 모두 있으면 존 씨 친구는 원하는 우유를 먹을 수 있는게 보장된다.
• 존씨 친구가 걷는 길에 한 종류의 소들만 있고, 해당 소가 존씨 친구 취향이 아니면 행복하지 않다.
• 즉, 같은 종류의 소들이 있는 농장끼리 연결되어 있으면 같은 그룹으로 묶는다.
• 존씨 친구가 가려는 시작점과 끝점이 다른 그룹이면 존씨 친구는 행복
• 같은 그룹이라면, 시작점 or 끝점이 존씨 친구 취향인지 확인 ( 풀이에는 시작점 기준으로 확인)

사용 알고리즘

농장이 tree 구조이고, 같은 종류 끼리 그룹으로 묶어야 한다.
그룹화를 할 수 있는 알고리즘에는 dfs or union-find 가 있는데
나는 union-find 를 사용했다.

풀이 방법

1. 먼저 입력 받은 정보를 바탕으로 union-find 알고리즘을 이용해 그룹화를 해준다.
2. 친구가 걷는 시작점과 끝점이 다른 그룹이면 1 출력
   소 종류가 달라고 다른 그룹이면, 중간에 다른 종류의 소가 있다는 뜻이므로 happy 하다.
3. 같은 그룹이면, 시작점의 소 종류를 확인하다. 친구가 선호하는 종류면 1 , 아니면 0 을 출력한다.

union - find 를 이용해 그룹화 하는 과정

위의 예시를 그림으로 표현하면 아래와 같은 트리 그래프가 만들어진다.

아직 그룹화가 되지 않은 상태며, 각 노드의 그룹장(parent) 는 자기 자신이다.
이를 배열로 표현하면 아래와 같다.

find

단순히 자신이 속한 tree 의 그룹장(parent) 를 찾으면 된다.

int find(int n){
    if(parent[n] == n) return n;
    return parent[n] = find(parent[n]);
}

자신이 이미 그룹장이면 본인을 return
아닐 경우, 재귀적으로 계속 위로 올라가 그룹장을 찾는다.
return 문에 parent[n] = find(parent[n]) 이렇게 해주면 그룹장을 찾은 후에 바로 본인의 parent 에 대입해주므로
다음 find 작업시 좀 더 빠르게 찾을 수 있다.

union

연결된 노드들을 입력 받으면서 그룹을 합치는 작업을(union or merge) 해보자

1) 1 , 2 입력
먼저 1번 농장과 2번 농장의 소가 같은 종류인지 확인한다.
여기서는 같은 종류이므로 merge 해준다.

merge 시에는 그냥 단순히 자신의 그룹장(parent) 로 다른 노드를 지정해주면 된다.
(나는 트리 높이가 높은 트리가 parent 로 가도록 최적화 작업을 넣어 주었다.)

여기서는 2번 노드의 parent 로 1번 노드를 정해주도록 하자.

2) 2 ,3 입력
2 번과 3번은 다른 종류의 소다. merge 작업을 하지 않고 그냥 넘어간다.

3) 2, 4 입력
같은 종류의 소다!

2번 노드의 parent 는 1번이고, 4번 노드의 parent 는 4다.
2번 노드의 tree의 높이가 더 크므로, 4번을 2번의 parent 인 1번에 붙여준다.

4) 1, 5 입력
다른 종류의 소다. 넘어간다.

농장은 이제 아래와 같은 그림같이 그룹화가 끝났다. (당연히 union-find 로 해서 만들어진 실제 tree 모양은 아래 같진 않다.)

코드

백준 통과한 코드다.
ios_base::sync_with_stdio(0); cin.tie(0); 을 넣어주지 않았을 때 시간초과가 발생했다.

#include <iostream>
#include <vector>
#include <string>
using namespace std;
/*
핵심 : 같은 것들끼리 그룹화 하자
*/
vector<int> parent;
vector<int> ranks; // tree 높이. merge 최적화 위해 사용
string cows;
//root 찾기
int find(int n){
    if(parent[n] == n) return n;
    return parent[n] = find(parent[n]);
}
//합치기
void merge(int a, int b) {
    // 같은 종류의 소가 아니라면 merge 하지 않는다.
    if(cows[a-1] != cows[b-1]) return;
    a = find(a); b = find(b);
    if(a==b) return;
    if(ranks[a] >= ranks[b]) { //높이가 낮은 트리가 높은 트리의 자식으로 간다.
        parent[b] = a;
    } else {
        parent[a] = b;
    }
    if(ranks[a] == ranks[b]){
        ranks[a]++;
    }

    return;
}
int main() {
    ios_base::sync_with_stdio(0); cin.tie(0); //이거 안해줬더니 시간초과남.뭐지 입력받는데 뭐 문제가 있나
    //scanf 로 하면 시간초과 안나나 뭐지
    int n, m;
    cin >> n >> m;
    cin >> cows;
    parent.reserve(n+1);
    ranks.resize(n+1,1);
    for(int i=0; i<=n; i++) {
        parent.push_back(i); //초기 그룹의 root 는 자기자신
    }
    for(int i=0; i<n-1; i++) {
        int a,b;
        cin >> a >> b;
        merge(a,b);
    }
    for(int i=0; i<m; i++) {
        int a,b;
        char c;
        cin >> a >> b >> c;
        if( find(a) != find(b) || cows[a-1] == c) 
            cout << "1"; //happy
        else
            cout << "0"; //unhappy
    }
    cout << endl;
}