[CF 1118F1] Tree Cutting (Easy Version)

[CF 1118F1] Tree Cutting (Easy Version) - 1800

문제 설명 : 

빨간색,파란색,무색(색깔 없음)노드로 이루어진 트리가 주어진다. 트리의 간선 (n-1)개 중에서 어느 한 간선을 제거했을 때 만들어지는 2개의 component가 있을텐데, 각 component를 이루는 노드의 색깔이 빨간색과 파란색이 동시에 존재하지 않도록 하는 간선이 몇 개인지 구하는 문제이다.

풀이 : 

처음 생각한 것은 각 간선별로 하나씩 지워보고 만들어지는 2개의 component의 노드 색깔을 조사해볼까 생각해봤는데, 노드 수(간선 수)가 최대 30만이므로 이 풀이로는 힘들 것 같았다. 그래서 어떻게든 한 번의 순회만으로 제거할 수 있는 간선을 찾는 것이 유력하다고 생각했다.

정답을 만족시키는 두 component에서 어느 한 component에 빨간색(혹은 파란색)만 존재한다는 의미는 곧 또 다른 component에는 상반된 색깔인 파란색(혹은 빨간색)만 존재한다고 해석할 수 있다. 또한 이 경우에는 특정 색의 상반된 색의 개수는 반드시 0개여야만 하낟. 그래서 나는 dfs를 돌면서 subtree를 이루는 색깔 노드의 개수를 root에게 계속 전달해 가다가, 어떤 subtree에서의 색깔의 개수가 전체 색깔의 총 개수와 일치하면서 상반된 색깔의 개수가 0개라면 해당 간선을 지우는 것이 가능한 것이다. 

주의할 점 : 

-

배울 점 : 

코드의 dfs부분에서 반환형이 pair<int,int>인데 first와 second는 각각 색깔의 수를 의미한다. dfs에서 parameter는 sub-dfs에게 전달해줄 내용을 갖고 있고, return type(pair<int,int>)는 sub-dfs가 root-dfs에게 전달해줄 내용을 갖고 있다.

아직 그래프문제는 같은 난이도의 타 유형에 비해서 취약한 것 같다. 앞으로 그래프 유형을 많이 풀어보면서 컨셉을 많이 접해보는 것이 중요할 것 같다. 다양한 컨셉을 접함으로써 그 유형의 다른 문제가 나왔을 때, 내가 취할 수 있는 연산이 어떤 것이 있는지 미리 대비해놓을 수 있기 때문이다.

코드 :

#include<stdio.h>
#include<math.h>
#include<string.h>
#include<iostream>
#include<functional>
#include<string>
#include<algorithm>
#include<vector>
#include<map>
#include<set>
#include<queue>
#include<assert.h>
#include<stdlib.h>
#include<stack>
using namespace std;
/*********************Contest Template***********************/
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
typedef pair<long long,long long> pll;
#define FASTIO ios_base::sync_with_stdio(false);cin.tie(0);
struct S{
    int a,b; S(){}S(int _a,int _b)
    { a=_a; b=_b; }
    const bool operator<(const S &o) const{
        return a<o.a;}
};
priority_queue<int,vector<int>,greater<int>> mpq;
const int SPACE=0,NL=1;  string exm;
inline void showAll(vector<int> &v,int sep){ //0=space,1="\n"
    for(int &here:v) printf("%d%c",here,(sep)?'\n':' '); }
inline void exf(void){ cout<<exm<<"\n"; exit(0); }
inline vector<int> int_seperation(int N,int d=10){
    vector<int> v; while(N){v.push_back(N%d); N/=d;}
    reverse(v.begin(),v.end()); return v; }
/*********************Contest Template***********************/
const int SIZE= 300009;
 
int arr[SIZE],color[3];
vector<int> tree[SIZE];
int ans=0;
 
pii dfs(int node,int par){
    int _1=(arr[node]==1),_2=(arr[node]==2);
    for(int next:tree[node]){
        if(par==next) continue;
        auto p=dfs(next,node);
        if(color[1]==p.first && p.second==0) ans++;
        if(color[2]==p.second && p.first==0) ans++;
        _1+=p.first,_2+=p.second;
    }
    return {_1,_2};
}
 
int main(){
    int n;    scanf("%d",&n);
 
    for(int i=1 ; i<=n ; i++){
        scanf("%d",&arr[i]);
        color[arr[i]]++;
    }
 
    for(int i=0 ; i<n-1 ; i++){
        int u,v;    scanf("%d %d",&u,&v);
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
 
    dfs(1,0);
    printf("%d\n",ans);
 
    return 0;
}
/*
 
*/