Problem p02986 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <cassert>
#include <cstdio>
#include <functional>
#include <tuple>
#include <vector>
using namespace std;
using int64 = long long int;

struct UnionFind {
  vector<int> node, anc, color;
  UnionFind(int n) : node(n, -1), anc(n), color(n) {}
  int find(int k) { return node[k] < 0 ? k : find(node[k]); }
  int size(int k) { return -node[find(k)]; }
  bool unite(int u, int v) {
    u = find(u), v = find(v);
    if (u == v)
      return false;
    if (node[u] < node[v])
      swap(u, v);
    node[u] += node[v];
    node[v] = u;
    return true;
  }
};

int main() {
  int N, Q;
  scanf("%d%d", &N, &Q);
  vector<vector<tuple<int, int, int>>> G(N);
  for (int i = 0; i < N - 1; i++) {
    int a, b, c, d;
    scanf("%d%d%d%d", &a, &b, &c, &d);
    a--;
    b--;
    c--;
    G[a].emplace_back(b, c, d);
    G[b].emplace_back(a, c, d);
  }

  vector<int64> ans(Q), neg(Q);
  vector<vector<tuple<int, int, int, int>>> info(N), tmp(N);
  for (int i = 0; i < Q; i++) {
    int x, y, u, v;
    scanf("%d%d%d%d", &x, &y, &u, &v);
    x--;
    u--;
    v--;
    info[u].emplace_back(x, y, i, v);
    info[v].emplace_back(x, y, i, u);
  }

  UnionFind uf(N);
  function<void(int, int)> find_lca = [&](int cur, int par) {
    uf.anc[cur] = cur;
    for (auto e : G[cur]) {
      int to = get<0>(e);
      if (to == par)
        continue;
      find_lca(to, cur);
      uf.unite(to, cur);
      uf.anc[uf.find(cur)] = cur;
    }
    uf.color[cur] = 1;
    for (auto e : info[cur]) {
      int col, d, idx, to;
      tie(col, d, idx, to) = e;
      if (to < 0)
        continue;
      if (uf.color[to] == 1 and neg[idx]++ == 0) {
        int lca = uf.anc[uf.find(to)];
        tmp[lca].emplace_back(col, d, idx, -1);
      }
    }
  };

  vector<int64> sum(N), cnt(N);
  int64 cost = 0;
  function<void(int, int)> go = [&](int cur, int par) {
    for (auto e : info[cur]) {
      int col, d, idx, p;
      tie(col, d, idx, p) = e;
      int64 diff = cost + cnt[col] * d - sum[col];
      if (p < 0) {
        ans[idx] -= 2 * diff;
      } else {
        ans[idx] += diff;
      }
    }

    for (auto e : G[cur]) {
      int to, col, dist;
      tie(to, col, dist) = e;
      if (to == par)
        continue;
      sum[col] += dist;
      cnt[col]++;
      cost += dist;
      go(to, cur);
      sum[col] -= dist;
      cnt[col]--;
      cost -= dist;
    }
  };

  find_lca(0, -1);
  for (int i = 0; i < N; i++)
    for (auto e : tmp[i])
      info[i].push_back(e);
  go(0, -1);
  for (auto a : ans)
    printf("%lld\n", a);
  return 0;
}


Fixed submission:
#include <algorithm>
#include <cassert>
#include <cstdio>
#include <functional>
#include <tuple>
#include <vector>
using namespace std;
using int64 = long long int;

struct UnionFind {
  vector<int> node, anc, color;
  UnionFind(int n) : node(n, -1), anc(n), color(n) {}
  int find(int k) { return node[k] < 0 ? k : find(node[k]); }
  int size(int k) { return -node[find(k)]; }
  bool unite(int u, int v) {
    u = find(u), v = find(v);
    if (u == v)
      return false;
    if (size(u) < size(v))
      swap(u, v);
    node[u] += node[v];
    node[v] = u;
    return true;
  }
};

int main() {
  int N, Q;
  scanf("%d%d", &N, &Q);
  vector<vector<tuple<int, int, int>>> G(N);
  for (int i = 0; i < N - 1; i++) {
    int a, b, c, d;
    scanf("%d%d%d%d", &a, &b, &c, &d);
    a--;
    b--;
    c--;
    G[a].emplace_back(b, c, d);
    G[b].emplace_back(a, c, d);
  }

  vector<int64> ans(Q), neg(Q);
  vector<vector<tuple<int, int, int, int>>> info(N), tmp(N);
  for (int i = 0; i < Q; i++) {
    int x, y, u, v;
    scanf("%d%d%d%d", &x, &y, &u, &v);
    x--;
    u--;
    v--;
    info[u].emplace_back(x, y, i, v);
    info[v].emplace_back(x, y, i, u);
  }

  UnionFind uf(N);
  function<void(int, int)> find_lca = [&](int cur, int par) {
    uf.anc[cur] = cur;
    for (auto e : G[cur]) {
      int to = get<0>(e);
      if (to == par)
        continue;
      find_lca(to, cur);
      uf.unite(to, cur);
      uf.anc[uf.find(cur)] = cur;
    }
    uf.color[cur] = 1;
    for (auto e : info[cur]) {
      int col, d, idx, to;
      tie(col, d, idx, to) = e;
      if (to < 0)
        continue;
      if (uf.color[to] == 1 and neg[idx]++ == 0) {
        int lca = uf.anc[uf.find(to)];
        tmp[lca].emplace_back(col, d, idx, -1);
      }
    }
  };

  vector<int64> sum(N), cnt(N);
  int64 cost = 0;
  function<void(int, int)> go = [&](int cur, int par) {
    for (auto e : info[cur]) {
      int col, d, idx, p;
      tie(col, d, idx, p) = e;
      int64 diff = cost + cnt[col] * d - sum[col];
      if (p < 0) {
        ans[idx] -= 2 * diff;
      } else {
        ans[idx] += diff;
      }
    }

    for (auto e : G[cur]) {
      int to, col, dist;
      tie(to, col, dist) = e;
      if (to == par)
        continue;
      sum[col] += dist;
      cnt[col]++;
      cost += dist;
      go(to, cur);
      sum[col] -= dist;
      cnt[col]--;
      cost -= dist;
    }
  };

  find_lca(0, -1);
  for (int i = 0; i < N; i++)
    for (auto e : tmp[i])
      info[i].push_back(e);
  go(0, -1);
  for (auto a : ans)
    printf("%lld\n", a);
  return 0;
}