Problem p01790 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <algorithm>
#include <cassert>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <ctime>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <string>
#include <vector>
using namespace std;

namespace {
using Integer = long long; //__int128;
template <class T> istream &operator>>(istream &is, vector<T> &vec) {
  for (T &val : vec)
    is >> val;
  return is;
}
template <class T> istream &operator,(istream &is, T &val) { return is >> val; }
template <class T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  for (int i = 0; i < vec.size(); i++)
    os << vec[i] << (i == vec.size() - 1 ? "" : " ");
  return os;
}
template <class T> ostream &operator,(ostream &os, const T &val) {
  return os << " " << val;
}

template <class H> void print(const H &head) { cout << head; }
template <class H, class... T> void print(const H &head, const T &...tail) {
  cout << head << " ";
  print(tail...);
}
template <class... T> void println(const T &...values) {
  print(values...);
  cout << endl;
}

template <class H> void eprint(const H &head) { cerr << head; }
template <class H, class... T> void eprint(const H &head, const T &...tail) {
  cerr << head << " ";
  eprint(tail...);
}
template <class... T> void eprintln(const T &...values) {
  eprint(values...);
  cerr << endl;
}

class range {
  long long start_, end_, step_;

public:
  struct range_iterator {
    long long val, step_;
    range_iterator(long long v, long long step) : val(v), step_(step) {}
    long long operator*() { return val; }
    void operator++() { val += step_; }
    bool operator!=(range_iterator &x) {
      return step_ > 0 ? val < x.val : val > x.val;
    }
  };
  range(long long len) : start_(0), end_(len), step_(1) {}
  range(long long start, long long end) : start_(start), end_(end), step_(1) {}
  range(long long start, long long end, long long step)
      : start_(start), end_(end), step_(step) {}
  range_iterator begin() { return range_iterator(start_, step_); }
  range_iterator end() { return range_iterator(end_, step_); }
};

string operator"" _s(const char *str, size_t size) { return move(string(str)); }
constexpr Integer my_pow(Integer x, Integer k, Integer z = 1) {
  return k == 0    ? z
         : k == 1  ? z * x
         : (k & 1) ? my_pow(x * x, k >> 1, z * x)
                   : my_pow(x * x, k >> 1, z);
}
constexpr Integer my_pow_mod(Integer x, Integer k, Integer M, Integer z = 1) {
  return k == 0    ? z % M
         : k == 1  ? z * x % M
         : (k & 1) ? my_pow_mod(x * x % M, k >> 1, M, z * x % M)
                   : my_pow_mod(x * x % M, k >> 1, M, z);
}
constexpr unsigned long long operator"" _ten(unsigned long long value) {
  return my_pow(10, value);
}

inline int k_bit(Integer x, int k) { return (x >> k) & 1; } // 0-indexed

mt19937 mt(chrono::duration_cast<chrono::nanoseconds>(
               chrono::steady_clock::now().time_since_epoch())
               .count());

template <class T> string join(const vector<T> &v, const string &sep) {
  stringstream ss;
  for (int i = 0; i < v.size(); i++) {
    if (i > 0)
      ss << sep;
    ss << v[i];
  }
  return ss.str();
}

string operator*(string s, int k) {
  string ret;
  while (k) {
    if (k & 1)
      ret += s;
    s += s;
    k >>= 1;
  }
  return ret;
}

} // namespace

constexpr long long mod = 9_ten + 7;

// O(v log v)
class centroid_decomposition {
public:
  const vector<vector<int>> &G;
  vector<vector<int>> H;
  int root;

  centroid_decomposition(const vector<vector<int>> &G_)
      : G(G_), H(G_.size()), subtree(G_.size(), 0), root(-1) {
    vector<bool> used(G.size(), false);
    rec(used, 0);
  }

private:
  vector<int> subtree;
  int get_centroid(vector<bool> &used, int p) {
    function<int(int, int)> dfs = [&](int pos, int last) {
      int ret = 1;
      for (int nx : G[pos]) {
        if (nx == last || used[nx])
          continue;
        ret += dfs(nx, pos);
      }
      return subtree[pos] = ret;
    };

    int num_v = dfs(p, p);

    function<int(int, int)> dfs2 = [&](int pos, int last) {
      pair<int, int> mx = {0, -1};
      for (int nx : G[pos]) {
        if (nx == last || used[nx])
          continue;
        mx = max(mx, pair<int, int>{subtree[nx], nx});
      }
      if (mx.first <= num_v / 2)
        return pos;
      return dfs2(mx.second, pos);
    };

    return dfs2(p, p);
  }

  int rec(vector<bool> &used, int p) {
    int r = get_centroid(used, p);
    if (root == -1)
      root = r;
    used[r] = true;
    for (int nx : G[r]) {
      if (!used[nx]) {
        H[r].push_back(rec(used, nx));
      }
    }
    return r;
  }
};

vector<bool> used;

long long get_count(vector<pair<int, int>> &out, vector<pair<int, int>> &in) {
  long long ret = 0;
  sort(in.begin(), in.end());
  sort(out.begin(), out.end());
  for (auto p : in) {
    ret +=
        lower_bound(out.begin(), out.end(),
                    pair<int, int>{-p.first + 1, -10000000}) -
        lower_bound(out.begin(), out.end(), pair<int, int>{-p.first, -p.first});
  }
  return ret;
}

void dfs(string &s, vector<vector<int>> &G, vector<pair<int, int>> &go_,
         vector<pair<int, int>> &come_, int pos, int par, pair<int, int> vo,
         pair<int, int> vi) {
  int d = (s[pos] == '(' ? 1 : -1);
  pair<int, int> vo_ = vo;
  vo_.first = vo.first + d;
  vo_.second = min(vo.second, vo_.first);

  pair<int, int> vi_ = vo;
  vi_.first = vi.first + d;
  vi_.second = min(0, vi.second + d);

  if (s[pos] == ')' && vo_.first <= 0)
    go_.push_back(vo_);
  if (vi_.second >= 0)
    come_.push_back(vi_);
  for (int nx : G[pos]) {
    if (nx == par || used[nx])
      continue;
    dfs(s, G, go_, come_, nx, pos, vo_, vi_);
  }
}

long long solve(vector<vector<int>> &G, vector<vector<int>> &H, string &s,
                int pos, int last) {
  used[pos] = true;

  long long ans = 0;
  for (int nx : H[pos]) {
    ans += solve(G, H, s, nx, pos);
  }

  //{sum, min}
  vector<pair<int, int>> go;   // pos -> v
  vector<pair<int, int>> come; // v -> pos

  for (int nx : G[pos]) {
    if (nx == last || used[nx])
      continue;

    vector<pair<int, int>> go_;   // pos -> v
    vector<pair<int, int>> come_; // v -> pos

    dfs(s, G, go_, come_, nx, pos, {0, 0},
        {s[pos] == '(' ? 1 : -1, s[pos] == '(' ? 0 : -1});
    ans -= get_count(go_, come_);
    go.insert(go.end(), go_.begin(), go_.end());
    come.insert(come.end(), come_.begin(), come_.end());
  }

  if (s[pos] == ')')
    for (auto p : come) {
      if (p.first == 0 && p.second >= 0)
        ans++;
    }
  if (s[pos] == '(')
    for (auto p : go) {
      if (p.first == -1 && p.second + 1 >= 0)
        ans++;
    }

  ans += get_count(go, come);

  used[pos] = false;

  return ans;
}

#include <cassert>
int main() {
  int n;
  cin >> n;
  string s;
  cin >> s;
  vector<vector<int>> G(n);
  for (int i : range(n - 1)) {
    int u, v;
    // cin >> u,v;
    scanf("%d%d", &u, &v);
    u--;
    v--;
    G[u].push_back(v);
    G[v].push_back(u);
  }

  centroid_decomposition cd(G);
  auto &H = cd.H;

  used.resize(n, false);
  long long ans = solve(G, H, s, cd.root, -1);

  println(ans);
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <ctime>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <string>
#include <vector>
using namespace std;

namespace {
using Integer = long long; //__int128;
template <class T> istream &operator>>(istream &is, vector<T> &vec) {
  for (T &val : vec)
    is >> val;
  return is;
}
template <class T> istream &operator,(istream &is, T &val) { return is >> val; }
template <class T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  for (int i = 0; i < vec.size(); i++)
    os << vec[i] << (i == vec.size() - 1 ? "" : " ");
  return os;
}
template <class T> ostream &operator,(ostream &os, const T &val) {
  return os << " " << val;
}

template <class H> void print(const H &head) { cout << head; }
template <class H, class... T> void print(const H &head, const T &...tail) {
  cout << head << " ";
  print(tail...);
}
template <class... T> void println(const T &...values) {
  print(values...);
  cout << endl;
}

template <class H> void eprint(const H &head) { cerr << head; }
template <class H, class... T> void eprint(const H &head, const T &...tail) {
  cerr << head << " ";
  eprint(tail...);
}
template <class... T> void eprintln(const T &...values) {
  eprint(values...);
  cerr << endl;
}

class range {
  long long start_, end_, step_;

public:
  struct range_iterator {
    long long val, step_;
    range_iterator(long long v, long long step) : val(v), step_(step) {}
    long long operator*() { return val; }
    void operator++() { val += step_; }
    bool operator!=(range_iterator &x) {
      return step_ > 0 ? val < x.val : val > x.val;
    }
  };
  range(long long len) : start_(0), end_(len), step_(1) {}
  range(long long start, long long end) : start_(start), end_(end), step_(1) {}
  range(long long start, long long end, long long step)
      : start_(start), end_(end), step_(step) {}
  range_iterator begin() { return range_iterator(start_, step_); }
  range_iterator end() { return range_iterator(end_, step_); }
};

string operator"" _s(const char *str, size_t size) { return move(string(str)); }
constexpr Integer my_pow(Integer x, Integer k, Integer z = 1) {
  return k == 0    ? z
         : k == 1  ? z * x
         : (k & 1) ? my_pow(x * x, k >> 1, z * x)
                   : my_pow(x * x, k >> 1, z);
}
constexpr Integer my_pow_mod(Integer x, Integer k, Integer M, Integer z = 1) {
  return k == 0    ? z % M
         : k == 1  ? z * x % M
         : (k & 1) ? my_pow_mod(x * x % M, k >> 1, M, z * x % M)
                   : my_pow_mod(x * x % M, k >> 1, M, z);
}
constexpr unsigned long long operator"" _ten(unsigned long long value) {
  return my_pow(10, value);
}

inline int k_bit(Integer x, int k) { return (x >> k) & 1; } // 0-indexed

mt19937 mt(chrono::duration_cast<chrono::nanoseconds>(
               chrono::steady_clock::now().time_since_epoch())
               .count());

template <class T> string join(const vector<T> &v, const string &sep) {
  stringstream ss;
  for (int i = 0; i < v.size(); i++) {
    if (i > 0)
      ss << sep;
    ss << v[i];
  }
  return ss.str();
}

string operator*(string s, int k) {
  string ret;
  while (k) {
    if (k & 1)
      ret += s;
    s += s;
    k >>= 1;
  }
  return ret;
}

} // namespace

constexpr long long mod = 9_ten + 7;

// O(v log v)
class centroid_decomposition {
public:
  const vector<vector<int>> &G;
  vector<vector<int>> H;
  int root;

  centroid_decomposition(const vector<vector<int>> &G_)
      : G(G_), H(G_.size()), subtree(G_.size(), 0), root(-1) {
    vector<bool> used(G.size(), false);
    rec(used, 0);
  }

private:
  vector<int> subtree;
  int get_centroid(vector<bool> &used, int p) {
    function<int(int, int)> dfs = [&](int pos, int last) {
      int ret = 1;
      for (int nx : G[pos]) {
        if (nx == last || used[nx])
          continue;
        ret += dfs(nx, pos);
      }
      return subtree[pos] = ret;
    };

    int num_v = dfs(p, p);

    function<int(int, int)> dfs2 = [&](int pos, int last) {
      pair<int, int> mx = {0, -1};
      for (int nx : G[pos]) {
        if (nx == last || used[nx])
          continue;
        mx = max(mx, pair<int, int>{subtree[nx], nx});
      }
      if (mx.first <= num_v / 2)
        return pos;
      return dfs2(mx.second, pos);
    };

    return dfs2(p, p);
  }

  int rec(vector<bool> &used, int p) {
    int r = get_centroid(used, p);
    if (root == -1)
      root = r;
    used[r] = true;
    for (int nx : G[r]) {
      if (!used[nx]) {
        H[r].push_back(rec(used, nx));
      }
    }
    return r;
  }
};

int main() {
  int n;
  cin >> n;
  string s;
  cin >> s;
  vector<vector<int>> G(n);
  for (int i : range(n - 1)) {
    int u, v;
    cin >> u, v;
    u--;
    v--;
    G[u].push_back(v);
    G[v].push_back(u);
  }

  centroid_decomposition cd(G);
  auto &H = cd.H;

  long long ans = 0;
  function<void(int, int)> solve = [&](int pos, int last) {
    for (int nx : H[pos]) {
      solve(nx, pos);
    }

    //{sum, min}
    vector<pair<int, int>> go;   // pos -> v
    vector<pair<int, int>> come; // v -> pos

    vector<pair<int, int>> go_;   // pos -> v
    vector<pair<int, int>> come_; // v -> pos

    function<void(int, int, pair<int, int>, pair<int, int>)> dfs =
        [&](int pos, int par, pair<int, int> v_out, pair<int, int> v_in) {
          int d = (s[pos] == '(' ? 1 : -1);
          pair<int, int> v_out__ = v_out;
          v_out__.first = v_out.first + d;
          v_out__.second = min(v_out.second, v_out__.first);

          pair<int, int> v_in__ = v_out;
          v_in__.first = v_in.first + d;
          v_in__.second = min(0, v_in.second + d);

          if (s[pos] == ')')
            go_.push_back(v_out__);
          if (v_in__.second >= 0)
            come_.push_back(v_in__);
          for (int nx : G[pos]) {
            if (nx == par || nx == last)
              continue;

            dfs(nx, pos, v_out__, v_in__);
          }
        };

    function<long long(vector<pair<int, int>> &, vector<pair<int, int>> &)>
        get_count =
            [&](vector<pair<int, int>> &out, vector<pair<int, int>> &in) {
              long long ret = 0;
              sort(in.begin(), in.end());
              sort(out.begin(), out.end());
              for (auto p : in) {
                ret += lower_bound(out.begin(), out.end(),
                                   pair<int, int>{-p.first + 1, -10000000}) -
                       lower_bound(out.begin(), out.end(),
                                   pair<int, int>{-p.first, -p.first});
              }
              return ret;
            };

    for (int nx : G[pos]) {
      if (nx == last)
        continue;
      dfs(nx, pos, {0, 0}, {s[pos] == '(' ? 1 : -1, s[pos] == '(' ? 0 : -1});
      ans -= get_count(go_, come_);
      go.insert(go.end(), go_.begin(), go_.end());
      come.insert(come.end(), come_.begin(), come_.end());
      go_.clear();
      come_.clear();
    }

    if (s[pos] == ')')
      for (auto p : come) {
        if (p.first == 0 && p.second >= 0)
          ans++;
      }
    if (s[pos] == '(')
      for (auto p : go) {
        if (p.first == -1 && p.second + 1 >= 0)
          ans++;
      }

    ans += get_count(go, come);
  };

  solve(cd.root, -1);

  println(ans);
  return 0;
}