Problem p03058 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <climits>
#include <cmath>
#include <complex>
#include <cstdint>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;
typedef long long ll;
typedef vector<int> vi;
typedef pair<int, int> pii;
#define MP make_pair
#define PB push_back
#define inf 1000000007
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define all(x) (x).begin(), (x).end()

template <typename A, size_t N, typename T>
void Fill(A (&array)[N], const T &val) {
  std::fill((T *)array, (T *)(array + N), val);
}

template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T> class segtree {
private:
  int n, sz;
  vector<pair<T, int>> node;

public:
  void resize(vector<T> &v) {
    sz = (int)v.size();
    n = 1;
    while (n < sz) {
      n *= 2;
    }
    node.resize(2 * n);
    for (int i = 0; i < sz; i++) {
      node[i + n] = make_pair(v[i], i);
    }
    for (int i = n - 1; i >= 1; i--) {
      node[i] = min(node[2 * i], node[2 * i + 1]);
    }
  }
  void update(int k, T a) {
    node[k += n] = make_pair(a, k);
    while (k >>= 1) {
      node[k] = min(node[2 * k], node[2 * k + 1]);
    }
  }
  pair<T, int> query(int a, int b) {
    pair<T, int> res1 = make_pair(numeric_limits<T>::max(), -1);
    pair<T, int> res2 = make_pair(numeric_limits<T>::max(), -1);
    a += n, b += n;
    while (a != b) {
      if (a % 2)
        res1 = min(res1, node[a++]);
      if (b % 2)
        res2 = min(res2, node[--b]);
      a >>= 1, b >>= 1;
    }
    return min(res1, res2);
  }
};

class LCA {
public:
  int V;
  vector<vector<int>> G;
  vector<int> ord;
  int depth[2001];
  int id[2001];
  int lca[2001][2001];
  segtree<int> st;
  LCA(){};
  LCA(int node_size) : V(node_size), G(V) {
    for (int i = 0; i < V; i++) {
      id[i] = -1;
      depth[i] = 0;
    }
  }
  void add_edge(int from, int to) {
    G[from].push_back(to), G[to].push_back(from);
  }
  void dfs(int u, int p, int k) {
    id[u] = (int)ord.size();
    ord.push_back(u);
    depth[u] = k;
    for (int v : G[u]) {
      if (v != p) {
        dfs(v, u, k + 1);
        ord.push_back(u);
      }
    }
  }
  void build() {
    ord.reserve(2 * V - 2);
    for (int i = 0; i < V; i++) {
      if (id[i] < 0) {
        dfs(i, -1, 0);
      }
    }
    vector<int> stvec(2 * V - 2);
    for (int i = 0; i < 2 * V - 2; i++) {
      stvec[i] = depth[ord[i]];
    }
    st.resize(stvec);
  }
  void Solve() {
    for (int i = 0; i < V; i++) {
      for (int j = i; j < V; j++) {
        lca[i][j] =
            ord[st.query(min(id[i], id[j]), max(id[i], id[j]) + 1).second];
      }
    }
  }
  int solve(int u, int v) {
    if (u > v)
      swap(u, v);
    return lca[u][v];
  }
  int dist(int u, int v) {
    if (u > v)
      swap(u, v);
    int lcaa = lca[u][v];
    return depth[u] + depth[v] - 2 * depth[lcaa];
  }
};

int n, m;
vector<vector<int>> g;
pair<int, int> p[2001];
vector<vector<int>> h;
LCA lca;
void edgeinit(int id) {
  if (p[id].first == -1)
    return;
  if (p[id].first == p[id].second) {
    p[id] = MP(-1, -1);
    return;
  }
  if (id == m)
    return;
  for (int i = 0; i < m; i++) {
    if (i == id)
      continue;
    if (p[i].first == -1)
      continue;
    if (p[i].first == p[id].first && p[i].second == p[id].second) {
      p[i] = MP(-1, -1);
      continue;
    }
    if (p[i].first == p[i].second) {
      p[i] = MP(-1, -1);
    }
    if (p[i].first == p[id].first) {
      int a = lca.dist(p[i].first, p[i].second);
      int b = lca.dist(p[id].first, p[id].second);
      int c = lca.dist(p[i].second, p[id].second);
      if (a - b == c) {
        p[i] = MP(p[i].second, p[id].second);
        edgeinit(i);
      } else if (b - a == c) {
        p[id] = MP(p[i].second, p[id].second);
        edgeinit(id);
      }
    } else if (p[i].first == p[id].second) {
      int a = lca.dist(p[i].first, p[i].second);
      int b = lca.dist(p[id].first, p[id].second);
      int c = lca.dist(p[i].second, p[id].first);
      if (a - b == c) {
        p[i] = MP(p[i].second, p[id].first);
        edgeinit(i);
      } else if (b - a == c) {
        p[id] = MP(p[i].second, p[id].first);
        edgeinit(id);
      }
    } else if (p[i].second == p[id].first) {
      int a = lca.dist(p[i].first, p[i].second);
      int b = lca.dist(p[id].first, p[id].second);
      int c = lca.dist(p[i].first, p[id].second);
      if (a - b == c) {

        p[i] = MP(p[i].first, p[id].second);
        edgeinit(i);
      } else if (b - a == c) {
        p[id] = MP(p[i].first, p[id].second);
        edgeinit(id);
      }
    } else if (p[i].second == p[id].second) {
      int a = lca.dist(p[i].first, p[i].second);
      int b = lca.dist(p[id].first, p[id].second);
      int c = lca.dist(p[i].first, p[id].first);
      if (a - b == c) {

        p[i] = MP(p[i].first, p[id].first);
        edgeinit(i);
      } else if (b - a == c) {
        p[id] = MP(p[i].first, p[id].first);
        edgeinit(id);
      }
    }
  }
}
int calc(int id) {
  queue<pair<int, int>> q;
  q.push(MP(id, id));
  int s = 0;
  while (!q.empty()) {
    auto x = q.front();
    q.pop();
    int dd = x.first;
    int pp = x.second;
    for (auto y : h[dd]) {
      if (lca.dist(pp, y) == lca.dist(pp, dd) + lca.dist(dd, y)) {
        s++;
        q.push(MP(y, dd));
      }
    }
  }
  return s;
}
int main() {
  cin >> n >> m;
  lca = LCA(n);
  g.resize(n);
  rep(i, n - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    lca.add_edge(a, b);
    g[a].push_back(b);
    g[b].push_back(a);
  }
  rep(i, m) {
    int c, d;
    cin >> c >> d;
    c--;
    d--;
    p[i] = (MP(c, d));
  }
  lca.build();
  lca.Solve();
  rep(i, m) { edgeinit(i); }
  h.resize(n);
  rep(i, m) {
    if (p[i].first == -1)
      continue;
    h[p[i].first].push_back(p[i].second);
    h[p[i].second].push_back(p[i].first);
  }
  int res = 0;
  rep(i, n) {
    res += calc(i);
    // cerr << i << " " << res << endl;
  }
  cout << res / 2 << endl;
  return 0;
}

Fixed submission:
#include <algorithm>
#include <bitset>
#include <cassert>
#include <chrono>
#include <climits>
#include <cmath>
#include <complex>
#include <cstdint>
#include <cstring>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;
typedef long long ll;
typedef vector<int> vi;
typedef pair<int, int> pii;
#define MP make_pair
#define PB push_back
#define inf 1000000007
#define rep(i, n) for (int i = 0; i < (int)(n); ++i)
#define all(x) (x).begin(), (x).end()

template <typename A, size_t N, typename T>
void Fill(A (&array)[N], const T &val) {
  std::fill((T *)array, (T *)(array + N), val);
}

template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T> class segtree {
private:
  int n, sz;
  vector<pair<T, int>> node;

public:
  void resize(vector<T> &v) {
    sz = (int)v.size();
    n = 1;
    while (n < sz) {
      n *= 2;
    }
    node.resize(2 * n);
    for (int i = 0; i < sz; i++) {
      node[i + n] = make_pair(v[i], i);
    }
    for (int i = n - 1; i >= 1; i--) {
      node[i] = min(node[2 * i], node[2 * i + 1]);
    }
  }
  void update(int k, T a) {
    node[k += n] = make_pair(a, k);
    while (k >>= 1) {
      node[k] = min(node[2 * k], node[2 * k + 1]);
    }
  }
  pair<T, int> query(int a, int b) {
    pair<T, int> res1 = make_pair(numeric_limits<T>::max(), -1);
    pair<T, int> res2 = make_pair(numeric_limits<T>::max(), -1);
    a += n, b += n;
    while (a != b) {
      if (a % 2)
        res1 = min(res1, node[a++]);
      if (b % 2)
        res2 = min(res2, node[--b]);
      a >>= 1, b >>= 1;
    }
    return min(res1, res2);
  }
};

class LCA {
public:
  int V;
  vector<vector<int>> G;
  vector<int> ord;
  int depth[2001];
  int id[2001];
  int lca[2001][2001];
  segtree<int> st;
  LCA(){};
  LCA(int node_size) : V(node_size), G(V) {
    for (int i = 0; i < V; i++) {
      id[i] = -1;
      depth[i] = 0;
    }
  }
  void add_edge(int from, int to) {
    G[from].push_back(to), G[to].push_back(from);
  }
  void dfs(int u, int p, int k) {
    id[u] = (int)ord.size();
    ord.push_back(u);
    depth[u] = k;
    for (int v : G[u]) {
      if (v != p) {
        dfs(v, u, k + 1);
        ord.push_back(u);
      }
    }
  }
  void build() {
    ord.reserve(2 * V - 2);
    for (int i = 0; i < V; i++) {
      if (id[i] < 0) {
        dfs(i, -1, 0);
      }
    }
    vector<int> stvec(2 * V - 2);
    for (int i = 0; i < 2 * V - 2; i++) {
      stvec[i] = depth[ord[i]];
    }
    st.resize(stvec);
  }
  void Solve() {
    for (int i = 0; i < V; i++) {
      for (int j = i; j < V; j++) {
        lca[i][j] =
            ord[st.query(min(id[i], id[j]), max(id[i], id[j]) + 1).second];
      }
    }
  }
  int solve(int u, int v) {
    if (u > v)
      swap(u, v);
    return lca[u][v];
  }
  int dist(int u, int v) {
    if (u > v)
      swap(u, v);
    int lcaa = lca[u][v];
    return depth[u] + depth[v] - 2 * depth[lcaa];
  }
};

int n, m;
vector<vector<int>> g;
pair<int, int> p[2001];
vector<vector<int>> h;
LCA lca;
void edgeinit(int id) {
  if (p[id].first == -1)
    return;
  if (p[id].first == p[id].second) {
    p[id] = MP(-1, -1);
    return;
  }
  if (id == m)
    return;
  for (int i = 0; i < m; i++) {
    if (i == id)
      continue;
    if (p[i].first == -1)
      continue;
    if (p[i].first == p[id].first && p[i].second == p[id].second) {
      p[i] = MP(-1, -1);
      continue;
    }
    if (p[i].first == p[i].second) {
      p[i] = MP(-1, -1);
    }
    if (p[i].first == p[id].first) {
      int a = lca.dist(p[i].first, p[i].second);
      int b = lca.dist(p[id].first, p[id].second);
      int c = lca.dist(p[i].second, p[id].second);
      if (a - b == c) {
        p[i] = MP(p[i].second, p[id].second);
        edgeinit(i);
      } else if (b - a == c) {
        p[id] = MP(p[i].second, p[id].second);
        edgeinit(id);
      }
    } else if (p[i].first == p[id].second) {
      int a = lca.dist(p[i].first, p[i].second);
      int b = lca.dist(p[id].first, p[id].second);
      int c = lca.dist(p[i].second, p[id].first);
      if (a - b == c) {
        p[i] = MP(p[i].second, p[id].first);
        edgeinit(i);
      } else if (b - a == c) {
        p[id] = MP(p[i].second, p[id].first);
        edgeinit(id);
      }
    } else if (p[i].second == p[id].first) {
      int a = lca.dist(p[i].first, p[i].second);
      int b = lca.dist(p[id].first, p[id].second);
      int c = lca.dist(p[i].first, p[id].second);
      if (a - b == c) {

        p[i] = MP(p[i].first, p[id].second);
        edgeinit(i);
      } else if (b - a == c) {
        p[id] = MP(p[i].first, p[id].second);
        edgeinit(id);
      }
    } else if (p[i].second == p[id].second) {
      int a = lca.dist(p[i].first, p[i].second);
      int b = lca.dist(p[id].first, p[id].second);
      int c = lca.dist(p[i].first, p[id].first);
      if (a - b == c) {

        p[i] = MP(p[i].first, p[id].first);
        edgeinit(i);
      } else if (b - a == c) {
        p[id] = MP(p[i].first, p[id].first);
        edgeinit(id);
      }
    }
  }
}
int calc(int id) {
  queue<pair<int, int>> q;
  q.push(MP(id, id));
  int s = 0;
  while (!q.empty()) {
    auto x = q.front();
    q.pop();
    int dd = x.first;
    int pp = x.second;
    for (auto y : h[dd]) {
      if (lca.dist(pp, y) == lca.dist(pp, dd) + lca.dist(dd, y)) {
        s++;
        q.push(MP(y, dd));
      }
    }
  }
  return s;
}
int main() {
  cin >> n >> m;
  lca = LCA(n);
  g.resize(n);
  rep(i, n - 1) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    lca.add_edge(a, b);
    g[a].push_back(b);
    g[b].push_back(a);
  }
  rep(i, m) {
    int c, d;
    cin >> c >> d;
    c--;
    d--;
    p[i] = (MP(c, d));
  }
  lca.build();
  lca.Solve();
  rep(i, m) { edgeinit(m - 1 - i); }
  h.resize(n);
  rep(i, m) {
    if (p[i].first == -1)
      continue;
    h[p[i].first].push_back(p[i].second);
    h[p[i].second].push_back(p[i].first);
  }
  int res = 0;
  rep(i, n) {
    res += calc(i);
    // cerr << i << " " << res << endl;
  }
  cout << res / 2 << endl;
  return 0;
}