Problem p02986 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <array>
#include <assert.h>
#include <bitset>
#include <cctype>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdlib>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <utility>

using namespace std;
typedef long long ll;

template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> T gcd(T a, T b) {
  a = abs(a), b = abs(b);
  while (b > 0) {
    tie(a, b) = make_pair(b, a % b);
  }
  return a;
}
// mt19937 rnd(chrono::steady_clock::now().time_since_epoch().count());

constexpr long long INF = 1LL << 60;
constexpr int inf = 1000000007;
constexpr long long mod = 1000000007LL;
// constexpr long long mod = 998244353LL;

struct HLD {
  vector<int> sz, in, out, head, par;
  vector<vector<int>> g;
  HLD() {}
  HLD(vector<vector<int>> &_g)
      : g(_g), sz(_g.size()), in(_g.size()), out(_g.size()), head(_g.size()),
        par(_g.size()) {}
  void dfs_sz(int cur, int pre) {
    sz[cur] = 1;
    for (int &next : g[cur]) {
      if (next == pre)
        continue;
      dfs_sz(next, cur);
      par[next] = cur;
      sz[cur] += sz[next];
      if (sz[next] > sz[g[cur][0]])
        swap(next, g[cur][0]);
    }
  }

  void dfs_hld(int cur, int pre, int &idx) {
    in[cur] = idx++;
    for (int &next : g[cur]) {
      if (next == pre)
        continue;
      if (next == g[cur][0])
        head[next] = head[cur];
      else
        head[next] = next;
      dfs_hld(next, cur, idx);
    }
    out[cur] = idx;
  }

  void build() {
    dfs_sz(0, -1);
    int idx = 0;
    dfs_hld(0, -1, idx);
  }

  void build(vector<vector<int>> &_g) {
    g = _g;
    sz.resize(g.size());
    in.resize(g.size());
    out.resize(g.size());
    head.resize(g.size());
    par.resize(g.size());

    dfs_sz(0, -1);
    int idx = 0;
    dfs_hld(0, -1, idx);
  }
  int lca(int u, int v) {
    for (;; v = par[head[v]]) {
      if (in[u] > in[v])
        swap(u, v);
      if (head[u] == head[v])
        return u;
    }
  }
};

template <typename T> struct BIT {
  int n;
  vector<T> bit;
  BIT() : n(0) {}
  BIT(int _n) : n(_n) { bit = vector<T>(n + 1); }
  void add1(int idx, T val) {
    for (int i = idx; i <= n; i += i & -i)
      bit[i] += val;
  }
  T sum1(int idx) {
    T res = 0;
    for (int i = idx; i > 0; i -= i & -i)
      res += bit[i];
    return res;
  }
  // 0-indexed
  void add(int idx, T val) { add1(idx + 1, val); }
  // 0-indexed [left, right)
  T sum(int left, int right) { return sum1(right) - sum1(left); }

  int lower_bound(T x) {
    int res = 0;
    int k = 1;
    while (2 * k <= n)
      k <<= 1;
    for (; k > 0; k >>= 1) {
      if (res + k <= n and bit[res + k] < x) {
        x -= bit[res + k];
        res += k;
      }
    }
    return res;
  }
};

int main() {

  cin.tie(nullptr);
  ios::sync_with_stdio(false);

  int n, Q;
  cin >> n >> Q;
  vector<vector<int>> g(n);
  vector<vector<tuple<int, int, int>>> vt(n);

  for (int i = 0; i < n - 1; i++) {
    int a, b, c, d;
    cin >> a >> b >> c >> d;
    a--;
    b--;
    c--;
    g[a].emplace_back(b);
    g[b].emplace_back(a);
    vt[c].emplace_back(a, b, d);
  }

  HLD hld(g);
  hld.build();
  // q[x] = y, u, v, id
  vector<vector<tuple<int, int, int, int>>> query(n);

  for (int i = 0; i < Q; i++) {
    int x, y, u, v;
    cin >> x >> y >> u >> v;
    x--;
    u--;
    v--;
    query[x].emplace_back(y, u, v, i);
  }
  vector<ll> res(n);
  BIT<ll> bit(n + 1), bc(n + 1);
  for (int i = 0; i < n; i++) {
    for (auto [u, v, d] : vt[i]) {
      if (hld.in[u] > hld.in[v])
        swap(u, v);
      bit.add(hld.in[v], d);
    }
  }
  for (int i = 0; i < n; i++) {
    for (auto [u, v, d] : vt[i]) {
      if (hld.in[u] > hld.in[v])
        swap(u, v);
      bit.add(hld.in[v], -d);
      bc.add(hld.in[v], 1);
    }
    for (auto [y, u, v, id] : query[i]) {
      ll ans = 0;
      ll cnt = 0;
      for (;; v = hld.par[hld.head[v]]) {
        if (hld.in[u] > hld.in[v])
          swap(u, v);
        if (hld.head[u] == hld.head[v]) {
          ans += bit.sum(hld.in[u] + 1, hld.in[v] + 1);
          cnt += bc.sum(hld.in[u] + 1, hld.in[v] + 1);
          break;
        } else {
          ans += bit.sum(hld.in[hld.head[v]], hld.in[v] + 1);
          cnt += bc.sum(hld.in[hld.head[v]], hld.in[v] + 1);
        }
      }
      ans += cnt * y;
      res[id] = ans;
    }
    for (auto [u, v, d] : vt[i]) {
      if (hld.in[u] > hld.in[v])
        swap(u, v);
      bit.add(hld.in[v], +d);
      bc.add(hld.in[v], -1);
    }
  }
  for (int i = 0; i < Q; i++)
    cout << res[i] << "\n";
  return 0;
}


Fixed submission:
#include <algorithm>
#include <array>
#include <assert.h>
#include <bitset>
#include <cctype>
#include <chrono>
#include <cmath>
#include <complex>
#include <cstdlib>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <utility>

using namespace std;
typedef long long ll;

template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return 1;
  }
  return 0;
}
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return 1;
  }
  return 0;
}
template <typename T> T gcd(T a, T b) {
  a = abs(a), b = abs(b);
  while (b > 0) {
    tie(a, b) = make_pair(b, a % b);
  }
  return a;
}
// mt19937 rnd(chrono::steady_clock::now().time_since_epoch().count());

constexpr long long INF = 1LL << 60;
constexpr int inf = 1000000007;
constexpr long long mod = 1000000007LL;
// constexpr long long mod = 998244353LL;

struct HLD {
  vector<int> sz, in, out, head, par;
  vector<vector<int>> g;
  HLD() {}
  HLD(vector<vector<int>> &_g)
      : g(_g), sz(_g.size()), in(_g.size()), out(_g.size()), head(_g.size()),
        par(_g.size()) {}
  void dfs_sz(int cur, int pre) {
    sz[cur] = 1;
    for (int &next : g[cur]) {
      if (next == pre)
        continue;
      dfs_sz(next, cur);
      par[next] = cur;
      sz[cur] += sz[next];
      if (sz[next] > sz[g[cur][0]])
        swap(next, g[cur][0]);
    }
  }

  void dfs_hld(int cur, int pre, int &idx) {
    in[cur] = idx++;
    for (int &next : g[cur]) {
      if (next == pre)
        continue;
      if (next == g[cur][0])
        head[next] = head[cur];
      else
        head[next] = next;
      dfs_hld(next, cur, idx);
    }
    out[cur] = idx;
  }

  void build() {
    dfs_sz(0, -1);
    int idx = 0;
    dfs_hld(0, -1, idx);
  }

  void build(vector<vector<int>> &_g) {
    g = _g;
    sz.resize(g.size());
    in.resize(g.size());
    out.resize(g.size());
    head.resize(g.size());
    par.resize(g.size());

    dfs_sz(0, -1);
    int idx = 0;
    dfs_hld(0, -1, idx);
  }
  int lca(int u, int v) {
    for (;; v = par[head[v]]) {
      if (in[u] > in[v])
        swap(u, v);
      if (head[u] == head[v])
        return u;
    }
  }
};

template <typename T> struct BIT {
  int n;
  vector<T> bit;
  BIT() : n(0) {}
  BIT(int _n) : n(_n) { bit = vector<T>(n + 1); }
  void add1(int idx, T val) {
    for (int i = idx; i <= n; i += i & -i)
      bit[i] += val;
  }
  T sum1(int idx) {
    T res = 0;
    for (int i = idx; i > 0; i -= i & -i)
      res += bit[i];
    return res;
  }
  // 0-indexed
  void add(int idx, T val) { add1(idx + 1, val); }
  // 0-indexed [left, right)
  T sum(int left, int right) { return sum1(right) - sum1(left); }

  int lower_bound(T x) {
    int res = 0;
    int k = 1;
    while (2 * k <= n)
      k <<= 1;
    for (; k > 0; k >>= 1) {
      if (res + k <= n and bit[res + k] < x) {
        x -= bit[res + k];
        res += k;
      }
    }
    return res;
  }
};

int main() {

  cin.tie(nullptr);
  ios::sync_with_stdio(false);

  int n, Q;
  cin >> n >> Q;
  vector<vector<int>> g(n);
  vector<vector<tuple<int, int, int>>> vt(n);

  for (int i = 0; i < n - 1; i++) {
    int a, b, c, d;
    cin >> a >> b >> c >> d;
    a--;
    b--;
    c--;
    g[a].emplace_back(b);
    g[b].emplace_back(a);
    vt[c].emplace_back(a, b, d);
  }

  HLD hld(g);
  hld.build();
  // q[x] = y, u, v, id
  vector<vector<tuple<int, int, int, int>>> query(n);

  for (int i = 0; i < Q; i++) {
    int x, y, u, v;
    cin >> x >> y >> u >> v;
    x--;
    u--;
    v--;
    query[x].emplace_back(y, u, v, i);
  }
  vector<ll> res(Q);
  BIT<ll> bit(n + 1), bc(n + 1);
  for (int i = 0; i < n; i++) {
    for (auto [u, v, d] : vt[i]) {
      if (hld.in[u] > hld.in[v])
        swap(u, v);
      bit.add(hld.in[v], d);
    }
  }
  for (int i = 0; i < n; i++) {
    for (auto [u, v, d] : vt[i]) {
      if (hld.in[u] > hld.in[v])
        swap(u, v);
      bit.add(hld.in[v], -d);
      bc.add(hld.in[v], 1);
    }
    for (auto [y, u, v, id] : query[i]) {
      ll ans = 0;
      ll cnt = 0;
      for (;; v = hld.par[hld.head[v]]) {
        if (hld.in[u] > hld.in[v])
          swap(u, v);
        if (hld.head[u] == hld.head[v]) {
          ans += bit.sum(hld.in[u] + 1, hld.in[v] + 1);
          cnt += bc.sum(hld.in[u] + 1, hld.in[v] + 1);
          break;
        } else {
          ans += bit.sum(hld.in[hld.head[v]], hld.in[v] + 1);
          cnt += bc.sum(hld.in[hld.head[v]], hld.in[v] + 1);
        }
      }
      ans += cnt * y;
      res[id] = ans;
    }
    for (auto [u, v, d] : vt[i]) {
      if (hld.in[u] > hld.in[v])
        swap(u, v);
      bit.add(hld.in[v], +d);
      bc.add(hld.in[v], -1);
    }
  }
  for (int i = 0; i < Q; i++)
    cout << res[i] << "\n";
  return 0;
}