Submission #761622

#	Time	Username	Problem	Language	Result	Execution time	Memory
761622		gun_gan	Zagrade (COI17_zagrade)	C++17	100 / 100	241 ms	124560 KiB

zagrade

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
 
const int MX = 3e5 + 5, inf = 1e9;
 
int N;
string s;
vector<int> g[MX];
 
int par[MX], bal[MX];
set<pair<int, int>> path[MX][2]; // open, close
// (balance, count)
 
int find(int v) {
      return par[v] == v ? v : par[v] = find(par[v]);
}
 
ll ans = 0;
 
void dfs(int v, int p) {
      for(auto u : g[v]) {
            if(u == p) continue;
            // build 
            bal[u] = bal[v] + (s[u] == '(' ? 1 : -1);
            dfs(u, v);
      }
 
      for(auto u : g[v]) {
            if(u == p) continue; 
 
            // erase invalid
            if(s[v] == ')') {
                  int pu = find(u);
                  auto it = path[pu][0].lower_bound({bal[v], -inf});

                  if(it != path[pu][0].end() && it->first == bal[v]) {
                  
                        path[pu][0].erase(it);
                  }
            }
      
            if(s[v] == '(') {
                  int pu = find(u);

                  auto it = path[pu][1].lower_bound({bal[v], -inf});
                  if(it != path[pu][1].end() && it->first == bal[v]) {
                        path[pu][1].erase(it);
                  } 
            }
      }
 
      // straight path
      if(s[v] == ')') {
            for(auto u : g[v]) {
                  if(u == p) continue;
 
                  int pu = find(u);
                  auto it = path[pu][0].lower_bound({bal[p], -inf});
                  if(it != path[pu][0].end() && it->first == bal[p]) {
                        ans += it->second;
                  }
            }
      } else {
            for(auto u : g[v]) {
                  if(u == p) continue;
 
                  int pu = find(u);
                  auto it = path[pu][1].lower_bound({bal[p], -inf});
                  if(it != path[pu][1].end() && it->first == bal[p]) {
                        ans += it->second;
                  }
            }
      }
 
      // calc answer, merging
      for(auto u : g[v]) {
            if(u == p) continue;
            int pu = find(u), pv = find(v);
 
 
            if((path[pu][0].size() + path[pu][1].size()) > (path[pv][0].size() + path[pv][1].size())) {
                  swap(pu, pv);
            }
 
            for(auto [b, c] : path[pu][0]) {
                  int k = b - bal[p];
                  k = -k + bal[v];
                  auto it = path[pv][1].lower_bound({k, -inf});
                  if(it != path[pv][1].end() && it->first == k) {
                        ans += 1LL * c * it->second;
                  }
            }
 
            for(auto [b, c] : path[pu][1]) {
                  int k = b - bal[p];
                  k = -k + bal[v];
                  auto it = path[pv][0].lower_bound({k, -inf});
                  if(it != path[pv][0].end() && it->first == k) {
                        ans += 1LL * c * it->second;
                  }
            }
 
 
            par[pu] = pv;
            for(auto x : path[pu][0]) {
                  auto it = path[pv][0].lower_bound({x.first, -inf});

                  if(it != path[pv][0].end() && it->first == x.first) {
                        path[pv][0].insert({x.first, x.second + it->second});
                        path[pv][0].erase(it);
                  } else {
                        path[pv][0].insert(x);
                  }
            }

             for(auto x : path[pu][1]) {
                  auto it = path[pv][1].lower_bound({x.first, -inf});
                  if(it != path[pv][1].end() && it->first == x.first) {
                        path[pv][1].insert({x.first, x.second + it->second});
                        path[pv][1].erase(it);
                  } else {
                        path[pv][1].insert(x);
                  }
            }

            path[pu][0].clear();
            path[pu][1].clear();
      }

      // add new path
      if(s[v] == '(') {
            int pv = find(v);
            auto it = path[pv][0].lower_bound({bal[v], -inf});

            if(it != path[pv][0].end() && it->first == bal[v]) {
                  path[pv][0].insert({it->first, it->second + 1});
                  path[pv][0].erase(*it);
            } else {
                  path[pv][0].insert({bal[v], 1});
            }
      } else {
            int pv = find(v);
            auto it = path[pv][1].lower_bound({bal[v], -inf});

            if(it != path[pv][1].end() && it->first == bal[v]) {
                  path[pv][1].insert({it->first, it->second + 1});
                  path[pv][1].erase(*it);
            } else {
                  path[pv][1].insert({bal[v], 1});
            }
      }
}     
 
int main() {
      ios_base::sync_with_stdio(0); cin.tie(0);
 
      cin >> N >> s;
      s = '#' + s;
 
      for(int i = 0; i < N - 1; i++) {
            int u, v;
            cin >> u >> v;
            g[u].push_back(v);
            g[v].push_back(u);
      }
 
      for(int i = 1; i <= N; i++) par[i] = i;
 
      bal[1] = (s[1] == '(' ? 1 : -1);
      dfs(1, 0);
 
      cout << ans << '\n';
}

• Subtask 1: 10 / 10
• Subtask 2: 30 / 30
• Subtask 3: 60 / 60