Problem p02986 (C++)

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

// #define endl "\n"
typedef long long ll;
using namespace std;
const int N = 2e5 + 10;
int n;
vector<int> g[N];
vector<int> e[N];
int sz[N], in[N], out[N], top[N], t, son[N];
;
int dep[N], par[N];
struct {
  int u, v, w, c;
} edge[N];
int val[N]; // 本来是记录点权，此处改为记录点对应的边的编号
int wei[N];
int col[N];
struct T {
  int val;
  map<int, int> mp0;
  map<int, int> mp1;
} tree[N];

void build(int k, int l, int r) {
  if (l == r) {
    tree[k].val = wei[l];
    tree[k].mp0[col[l]]++;
    tree[k].mp1[col[l]] += wei[l];
    return;
  }
  int mid = (l + r) >> 1;
  build(k * 2, l, mid);
  build(k * 2 + 1, mid + 1, r);
  tree[k].val = tree[2 * k].val + tree[2 * k + 1].val;
  for (auto p : tree[2 * k].mp0)
    tree[k].mp0[p.first] += p.second;
  for (auto p : tree[2 * k + 1].mp0)
    tree[k].mp0[p.first] += p.second;
  for (auto p : tree[2 * k].mp1)
    tree[k].mp1[p.first] += p.second;
  for (auto p : tree[2 * k + 1].mp1)
    tree[k].mp1[p.first] += p.second;
}

struct P {
  int a, b, c;
};

P query(int k, int l, int r, int lef, int rig, int color) {
  if (lef <= l && r <= rig) {
    if (tree[k].mp0.count(color))
      return {tree[k].val, tree[k].mp0[color], tree[k].mp1[color]};
    else
      return {tree[k].val, 0, 0};
  }
  P ans = {0, 0, 0};
  int mid = (l + r) >> 1;
  if (lef <= mid) {
    auto t = query(k * 2, l, mid, lef, rig, color);
    ans.a += t.a;
    ans.b += t.b;
    ans.c += t.c;
  }
  if (rig > mid) {
    auto t = query(k * 2 + 1, mid + 1, r, lef, rig, color);
    ans.a += t.a;
    ans.b += t.b;
    ans.c += t.c;
  }
  return ans;
}

// 注意v是引用！因为要swap
void dfs_sz(int u, int from) {
  sz[u] = 1;
  son[u] = 0;
  for (int i = 0; i < g[u].size(); i++) {
    int &v = g[u][i], &w = e[u][i];
    if (v == from)
      continue;
    dep[v] = dep[u] + 1, par[v] = u;
    dfs_sz(v, u);
    sz[u] += sz[v];
    val[v] = w;
    if (sz[v] > sz[son[u]])
      son[u] = v;
  }
}

void dfs_hld(int u, int from) {
  in[u] = ++t;
  if (!son[u])
    return;
  top[son[u]] = top[u];
  dfs_hld(son[u], u);
  for (auto v : g[u]) {
    if (v == from || v == son[u])
      continue;
    top[v] = v;
    dfs_hld(v, u);
  }
  out[u] = t;
}

P querys(int x, int y, int color) {
  P s = {0, 0, 0};
  while (top[x] != top[y]) {
    if (dep[top[x]] < dep[top[y]])
      swap(x, y);
    // 此处把区间映射到[in[top[x]],in[x]]!
    auto p = query(1, 1, t, in[top[x]], in[x], color);
    s.a += p.a;
    s.b += p.b;
    s.c += p.c;
    x = par[top[x]];
  }
  if (in[x] > in[y])
    swap(x, y);
  // 此处把区间映射到[in[x],in[y]]!
  // 因为这道题是对边的权值求解，然后处理方式是把边权映射到点下面的点上
  // 所以最上面的点对应的边不在路径里，所以此处把区间映射到[in[x]+1,in[y]]
  auto p = query(1, 1, t, in[x] + 1, in[y], color);
  s.a += p.a;
  s.b += p.b;
  s.c += p.c;
  return s;
}

void init(int root) {
  top[root] = root, dep[root] = 1;
  dfs_sz(root, -1), dfs_hld(root, -1);
}

int q;

int main() {
  ios::sync_with_stdio(0), cin.tie(0);
  cin >> n >> q;
  for (int i = 1; i <= n - 1; i++) {
    int u, v, w, c;
    cin >> u >> v >> c >> w;
    g[u].push_back(v), g[v].push_back(u);
    e[u].push_back(i), e[v].push_back(i);
    edge[i] = {u, v, w, c};
  }
  init(1);
  for (int i = 2; i <= n; i++) {
    wei[in[i]] = edge[val[i]].w;
    col[in[i]] = edge[val[i]].c;
  }
  build(1, 1, t);
  while (q--) {
    int x, y, u, v;
    cin >> x >> y >> u >> v;
    auto p = querys(u, v, x);
    cout << p.a - p.c + p.b * y << endl;
  }
}

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

// #define endl "\n"
typedef long long ll;
using namespace std;
const int N = 2e5 + 10;
int n;
vector<int> g[N];
vector<int> e[N];
int sz[N], in[N], out[N], top[N], t, son[N];
;
int dep[N], par[N];
struct {
  int u, v, w, c;
} edge[N];
int val[N]; // 本来是记录点权，此处改为记录点对应的边的编号
int wei[N];
int col[N];
struct T {
  int val;
  map<int, int> mp0;
  map<int, int> mp1;
} tree[N * 4];

void build(int k, int l, int r) {
  if (l == r) {
    tree[k].val = wei[l];
    tree[k].mp0[col[l]]++;
    tree[k].mp1[col[l]] += wei[l];
    return;
  }
  int mid = (l + r) >> 1;
  build(k * 2, l, mid);
  build(k * 2 + 1, mid + 1, r);
  tree[k].val = tree[2 * k].val + tree[2 * k + 1].val;
  for (auto p : tree[2 * k].mp0)
    tree[k].mp0[p.first] += p.second;
  for (auto p : tree[2 * k + 1].mp0)
    tree[k].mp0[p.first] += p.second;
  for (auto p : tree[2 * k].mp1)
    tree[k].mp1[p.first] += p.second;
  for (auto p : tree[2 * k + 1].mp1)
    tree[k].mp1[p.first] += p.second;
}

struct P {
  int a, b, c;
};

P query(int k, int l, int r, int lef, int rig, int color) {
  if (lef <= l && r <= rig) {
    if (tree[k].mp0.count(color))
      return {tree[k].val, tree[k].mp0[color], tree[k].mp1[color]};
    else
      return {tree[k].val, 0, 0};
  }
  P ans = {0, 0, 0};
  int mid = (l + r) >> 1;
  if (lef <= mid) {
    auto t = query(k * 2, l, mid, lef, rig, color);
    ans.a += t.a;
    ans.b += t.b;
    ans.c += t.c;
  }
  if (rig > mid) {
    auto t = query(k * 2 + 1, mid + 1, r, lef, rig, color);
    ans.a += t.a;
    ans.b += t.b;
    ans.c += t.c;
  }
  return ans;
}

// 注意v是引用！因为要swap
void dfs_sz(int u, int from) {
  sz[u] = 1;
  son[u] = 0;
  for (int i = 0; i < g[u].size(); i++) {
    int &v = g[u][i], &w = e[u][i];
    if (v == from)
      continue;
    dep[v] = dep[u] + 1, par[v] = u;
    dfs_sz(v, u);
    sz[u] += sz[v];
    val[v] = w;
    if (sz[v] > sz[son[u]])
      son[u] = v;
  }
}

void dfs_hld(int u, int from) {
  in[u] = ++t;
  if (!son[u])
    return;
  top[son[u]] = top[u];
  dfs_hld(son[u], u);
  for (auto v : g[u]) {
    if (v == from || v == son[u])
      continue;
    top[v] = v;
    dfs_hld(v, u);
  }
  out[u] = t;
}

P querys(int x, int y, int color) {
  P s = {0, 0, 0};
  while (top[x] != top[y]) {
    if (dep[top[x]] < dep[top[y]])
      swap(x, y);
    // 此处把区间映射到[in[top[x]],in[x]]!
    auto p = query(1, 1, t, in[top[x]], in[x], color);
    s.a += p.a;
    s.b += p.b;
    s.c += p.c;
    x = par[top[x]];
  }
  if (in[x] > in[y])
    swap(x, y);
  // 此处把区间映射到[in[x],in[y]]!
  // 因为这道题是对边的权值求解，然后处理方式是把边权映射到点下面的点上
  // 所以最上面的点对应的边不在路径里，所以此处把区间映射到[in[x]+1,in[y]]
  auto p = query(1, 1, t, in[x] + 1, in[y], color);
  s.a += p.a;
  s.b += p.b;
  s.c += p.c;
  return s;
}

void init(int root) {
  top[root] = root, dep[root] = 1;
  dfs_sz(root, -1), dfs_hld(root, -1);
}

int q;

int main() {
  ios::sync_with_stdio(0), cin.tie(0);
  cin >> n >> q;
  for (int i = 1; i <= n - 1; i++) {
    int u, v, w, c;
    cin >> u >> v >> c >> w;
    g[u].push_back(v), g[v].push_back(u);
    e[u].push_back(i), e[v].push_back(i);
    edge[i] = {u, v, w, c};
  }
  init(1);
  for (int i = 2; i <= n; i++) {
    wei[in[i]] = edge[val[i]].w;
    col[in[i]] = edge[val[i]].c;
  }
  build(1, 1, t);
  while (q--) {
    int x, y, u, v;
    cin >> x >> y >> u >> v;
    auto p = querys(u, v, x);
    cout << p.a - p.c + p.b * y << endl;
  }
}