얼마나 예뻐? - CyberFlower Algorithm

2019 KCPC 얼마나 예뻐? (신입생 부문 H번 / 일반 부문 D번) 코드 제출 url

올바른 괄호 문자열을 조금 바꾸어 말하자면, ‘(‘을 +1로, ‘)’을 -1로 대응하였을 때 prefix sum의 최솟값이 0이고 total sum이 0인 문자열을 의미합니다.

( prefix sum이란 $ prefix_{sum}[i]=\sum_{0}^{i}arr[j] $ 으로, 0번부터 i번까지의 합을 의미합니다.)

따라서 우리는 어떤 sub tree가 pretty 한지 알기 위해서 자식 노드들의 (total_sum, min(prefix_sum) 만 알면 됩니다.

어떤 sub tree를 1(루트) $ \rightarrow $ 2 $ \rightarrow $ 3 $ \rightarrow $ 4 … 이런 식으로 진행된다면, x번째 자식 sub tree로 새로 생기는 최소값은 [1 … x-1] sum + min(x를 루트로 하는 subtree) 이므로 합을 구하는 과정에서 쉽게 같이 구할 수 있습니다.

#include <bits/stdc++.h>
#define ff first
#define ss second
#define pb push_back
using namespace std;
typedef pair<int,int> pii;
const int SZ=200005;
vector<int> graph[SZ];
int bra[SZ];
int tot;
pii tra(int p, int x){
    int sm=bra[x];
    int mn=bra[x];
    for(int nxt:graph[x]){
        if(nxt==p) continue;
        pii xx=tra(x,nxt);
        mn=min(mn,sm+xx.ss);
        sm+=xx.ff;
    }
    if(mn==0 && sm==0) tot++;
    return {sm,mn};
}
int main(void){
    ios_base::sync_with_stdio(false); cin.tie(NULL);
    int n; cin>>n;
    for(int i=1;i<=n;i++){
        int x; cin>>x;
        bra[i]=(x==0?1:-1);
    }
    for(int i=1;i<n;i++){
        int x,y; cin>>x>>y;
        graph[x].pb(y);
        graph[y].pb(x);
    }
    for(int i=1;i<=n;i++) sort(graph[i].begin(),graph[i].end());
    tra(0,1);
    cout<<tot;
    return 0;
}

2019 KCPC pretty