Submission #761627

#	Time	Username	Problem	Language	Result	Execution time	Memory
761627		gun_gan	Zagrade (COI17_zagrade)	C++17	100 / 100	236 ms	125176 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);
      }

      int c = s[v] == '(';
 
      for(auto u : g[v]) {
            if(u == p) continue; 
 
            // erase invalid
            int pu = find(u);
            auto it = path[pu][c].lower_bound({bal[v], -inf});

            if(it != path[pu][c].end() && it->first == bal[v]) {
                  path[pu][c].erase(it);
            }
      }
 
      // straight path
      for(auto u : g[v]) {
            if(u == p) continue;

            int pu = find(u);
            auto it = path[pu][c].lower_bound({bal[p], -inf});
            if(it != path[pu][c].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(int t = 0; t < 2; t++) {
                  // calc answer
                  for(auto [b, c] : path[pu][t]) {
                        int k = b - bal[p];
                        k = -k + bal[v];
                        auto it = path[pv][t ^ 1].lower_bound({k, -inf});
                        if(it != path[pv][t ^ 1].end() && it->first == k) {
                              ans += 1LL * c * it->second;
                        }
                  }
            }
            par[pu] = pv;

            for(int t = 0; t < 2; t++) {
                  // merging
                  for(auto x : path[pu][t]) {
                        auto it = path[pv][t].lower_bound({x.first, -inf});

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

      // add new path
      int pv = find(v);
      auto it = path[pv][c ^ 1].lower_bound({bal[v], -inf});

      if(it != path[pv][c ^ 1].end() && it->first == bal[v]) {
            path[pv][c ^ 1].insert({it->first, it->second + 1});
            path[pv][c ^ 1].erase(*it);
      } else {
            path[pv][c ^ 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