TOCPC2021: Royal Parentheses

#include <bits/stdc++.h>
using namespace std;

typedef long long lli;
typedef pair<int, int> pii;

const int N = 1e5 + 5;

vector<int> adj[N];
pii pttn[2][N]; // (stack count, min/max of stack count)
char str[N];

int main(){

    int nVertex, st, ed;
    scanf("%d%d%d", &nVertex, &st, &ed);
    for(int i = 1; i < nVertex; ++i){
        int u, v;
        scanf("%d%d", &u, &v);
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    scanf(" %s", str + 1);

    // BFS from st (Index 0 of pttn)
    for(int u = 1; u <= nVertex; ++u){
        pttn[0][u] = pii(1e9, 0);
    }
    queue<int> que;
    pttn[0][st] = str[st] == '(' ? pii(1, 0) : pii(-1, -1);
    que.push(st);
    while(!que.empty()){
        int u = que.front();
        que.pop();
        for(int v : adj[u]){
            if(pttn[0][v] == pii(1e9, 0)){
                if(str[v] == '('){
                    pttn[0][v].first = pttn[0][u].first + 1;
                } else {
                    pttn[0][v].first = pttn[0][u].first - 1;
                }
                pttn[0][v].second = min(pttn[0][u].second, pttn[0][v].first);
                que.push(v);
            }
        }
    }

    // BFS from ed (Index 1 of pttn)
    for(int u = 1; u <= nVertex; ++u){
        pttn[1][u] = pii(1e9, 0);
    }
    pttn[1][ed] = str[ed] == '(' ? pii(1, 1) : pii(-1, 0);
    que.push(ed);
    while(!que.empty()){
        int u = que.front();
        que.pop();
        for(int v : adj[u]){
            if(pttn[1][v] == pii(1e9, 0)){
                if(str[v] == '('){
                    pttn[1][v].first = pttn[1][u].first + 1;
                } else {
                    pttn[1][v].first = pttn[1][u].first - 1;
                }
                pttn[1][v].second = max(pttn[1][u].second, pttn[1][v].first);
                que.push(v);
            }
        }
    }

    // Count valid front parentheses
    map<int, int> mp;
    for(int u = 1; u <= nVertex; ++u){
        if(pttn[0][u].second == 0){
            ++mp[pttn[0][u].first];
        }
    }

    // Count valid pairs of parentheses
    lli cnt = 0;
    for(int u = 1; u <= nVertex; ++u){
        if(pttn[1][u].second == 0){
            map<int, int>::iterator itr = mp.find(-pttn[1][u].first);
            if(itr != mp.end()){
                cnt += itr -> second;
            }
        }
    }
    cout << cnt;

    return 0;
}