Problem p02986 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>

using namespace std;

struct disjoint_set {
  vector<int> p;

  disjoint_set(int n) : p(n, -1) {}

  int root(int u) { return p[u] < 0 ? u : p[u] = root(p[u]); }

  bool merge(int u, int v) {
    if ((u = root(u)) == (v = root(v)))
      return false;
    if (p[u] > p[v])
      swap(u, v);
    p[u] += p[v];
    p[v] = u;
    return true;
  }
};

vector<int> lca_offline(const vector<vector<int>> &tree, int root,
                        const vector<pair<int, int>> &queries) {
  const int n = tree.size();
  vector<vector<int>> queries_id(n);
  for (int i = 0; i < (int)queries.size(); ++i) {
    queries_id[queries[i].first].emplace_back(i);
    if (queries[i].first != queries[i].second)
      queries_id[queries[i].second].emplace_back(i);
  }
  vector<int> anc(n), color(n), ans(queries.size());
  disjoint_set ds(n);
  function<void(int)> dfs = [&](int u) {
    anc[u] = u;
    color[u] = -1;
    for (int v : tree[u])
      if (!color[v]) {
        dfs(v);
        ds.merge(u, v);
        anc[ds.root(u)] = u;
      }
    color[u] = 1;
    for (int q : queries_id[u]) {
      int other = queries[q].first + queries[q].second - u;
      if (color[other] == 1)
        ans[q] = anc[ds.root(other)];
    }
  };
  dfs(root);
  return ans;
}

int main() {
#ifdef jcg
  assert(freopen("input.in", "r", stdin));
//    assert(freopen("output.out", "w", stdout));
#endif

  ios_base::sync_with_stdio(0);
  cin.tie(0);

  int N, Q;
  cin >> N >> Q;

  vector<int> a(N - 1), b(N - 1), c(N - 1), d(N - 1);
  for (int i = 0; i < N - 1; ++i) {
    cin >> a[i] >> b[i] >> c[i] >> d[i];
    --a[i], --b[i];
  }

  vector<int> x(Q), y(Q), u(Q), v(Q);
  for (int i = 0; i < Q; ++i) {
    cin >> x[i] >> y[i] >> u[i] >> v[i];
    --u[i], --v[i];
  }

  vector<int> lca;
  {
    vector<vector<int>> g(N);
    vector<pair<int, int>> lca_queries;
    for (int i = 0; i < N - 1; ++i) {
      g[a[i]].emplace_back(b[i]);
      g[b[i]].emplace_back(a[i]);
    }
    for (int i = 0; i < Q; ++i)
      lca_queries.emplace_back(u[i], v[i]);
    lca = lca_offline(g, 0, lca_queries);
  }

  vector<pair<int, int>> queries;
  vector<vector<int>> queries_id(N);

  for (int i = 0; i < Q; ++i) {
    queries_id[u[i]].emplace_back(queries.size());
    queries.emplace_back(x[i], y[i]);

    queries_id[v[i]].emplace_back(queries.size());
    queries.emplace_back(x[i], y[i]);

    queries_id[lca[i]].emplace_back(queries.size());
    queries.emplace_back(x[i], y[i]);
  }

  vector<vector<tuple<int, int, int>>> tree(N);
  for (int i = 0; i < N - 1; ++i) {
    tree[a[i]].emplace_back(b[i], c[i], d[i]);
    tree[b[i]].emplace_back(a[i], c[i], d[i]);
  }

  vector<int> ans(Q), color_sum(N), freq(N);

  function<void(int, int, int)> dfs = [&](int u, int p, int d) {
    for (int q : queries_id[u]) {
      int xq, yq;
      tie(xq, yq) = queries[q];
      ans[q] = d + freq[xq] * yq - color_sum[xq];
    }
    for (auto e : tree[u]) {
      int v, xe, ye;
      tie(v, xe, ye) = e;
      if (v == p)
        continue;
      ++freq[xe];
      color_sum[xe] += ye;
      dfs(v, u, d + ye);
      --freq[xe];
      color_sum[xe] -= ye;
    }
  };

  dfs(0, 0, 0);

  for (int i = 0; i < 3 * Q; i += 3)
    cout << (ans[i] - ans[i + 2]) + (ans[i + 1] - ans[i + 2]) << endl;

  return 0;
}


Fixed submission:
#include <bits/stdc++.h>

using namespace std;

struct disjoint_set {
  vector<int> p;

  disjoint_set(int n) : p(n, -1) {}

  int root(int u) { return p[u] < 0 ? u : p[u] = root(p[u]); }

  bool merge(int u, int v) {
    if ((u = root(u)) == (v = root(v)))
      return false;
    if (p[u] > p[v])
      swap(u, v);
    p[u] += p[v];
    p[v] = u;
    return true;
  }
};

vector<int> lca_offline(const vector<vector<int>> &tree, int root,
                        const vector<pair<int, int>> &queries) {
  const int n = tree.size();
  vector<vector<int>> queries_id(n);
  for (int i = 0; i < (int)queries.size(); ++i) {
    queries_id[queries[i].first].emplace_back(i);
    if (queries[i].first != queries[i].second)
      queries_id[queries[i].second].emplace_back(i);
  }
  vector<int> anc(n), color(n), ans(queries.size());
  disjoint_set ds(n);
  function<void(int)> dfs = [&](int u) {
    anc[u] = u;
    color[u] = -1;
    for (int v : tree[u])
      if (!color[v]) {
        dfs(v);
        ds.merge(u, v);
        anc[ds.root(u)] = u;
      }
    color[u] = 1;
    for (int q : queries_id[u]) {
      int other = queries[q].first + queries[q].second - u;
      if (color[other] == 1)
        ans[q] = anc[ds.root(other)];
    }
  };
  dfs(root);
  return ans;
}

int main() {
#ifdef jcg
  assert(freopen("input.in", "r", stdin));
//    assert(freopen("output.out", "w", stdout));
#endif

  ios_base::sync_with_stdio(0);
  cin.tie(0);

  int N, Q;
  cin >> N >> Q;

  vector<int> a(N - 1), b(N - 1), c(N - 1), d(N - 1);
  for (int i = 0; i < N - 1; ++i) {
    cin >> a[i] >> b[i] >> c[i] >> d[i];
    --a[i], --b[i];
  }

  vector<int> x(Q), y(Q), u(Q), v(Q);
  for (int i = 0; i < Q; ++i) {
    cin >> x[i] >> y[i] >> u[i] >> v[i];
    --u[i], --v[i];
  }

  vector<int> lca;
  {
    vector<vector<int>> g(N);
    vector<pair<int, int>> lca_queries;
    for (int i = 0; i < N - 1; ++i) {
      g[a[i]].emplace_back(b[i]);
      g[b[i]].emplace_back(a[i]);
    }
    for (int i = 0; i < Q; ++i)
      lca_queries.emplace_back(u[i], v[i]);
    lca = lca_offline(g, 0, lca_queries);
  }

  vector<pair<int, int>> queries;
  vector<vector<int>> queries_id(N);

  for (int i = 0; i < Q; ++i) {
    queries_id[u[i]].emplace_back(queries.size());
    queries.emplace_back(x[i], y[i]);

    queries_id[v[i]].emplace_back(queries.size());
    queries.emplace_back(x[i], y[i]);

    queries_id[lca[i]].emplace_back(queries.size());
    queries.emplace_back(x[i], y[i]);
  }

  vector<vector<tuple<int, int, int>>> tree(N);
  for (int i = 0; i < N - 1; ++i) {
    tree[a[i]].emplace_back(b[i], c[i], d[i]);
    tree[b[i]].emplace_back(a[i], c[i], d[i]);
  }

  vector<int> ans(3 * Q), color_sum(N), freq(N);

  function<void(int, int, int)> dfs = [&](int u, int p, int d) {
    for (int q : queries_id[u]) {
      int xq, yq;
      tie(xq, yq) = queries[q];
      ans[q] = d + freq[xq] * yq - color_sum[xq];
    }
    for (auto e : tree[u]) {
      int v, xe, ye;
      tie(v, xe, ye) = e;
      if (v == p)
        continue;
      ++freq[xe];
      color_sum[xe] += ye;
      dfs(v, u, d + ye);
      --freq[xe];
      color_sum[xe] -= ye;
    }
  };

  dfs(0, 0, 0);

  for (int i = 0; i < 3 * Q; i += 3)
    cout << (ans[i] - ans[i + 2]) + (ans[i + 1] - ans[i + 2]) << endl;

  return 0;
}