Problem p02986 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <iostream>
#include <stdio.h>
#include <string.h>
#include <vector>
using namespace std;

const int M = 1e5 + 5;
const int MB = 19;

int n, q;

struct TN {
  int ls, rs;
  int cnt, sum;
  TN(int ls = 0, int rs = 0, int cnt = 0, int sum = 0)
      : ls(ls), rs(rs), cnt(cnt), sum(sum) {}
};

struct PSegT {
  TN tn[M * MB];
  int nn; // number of node
  int P, V;
  void init(int n) {
    nn = 0;
    int root;
    build(root, 0, n);
  }
  void build(int &o, int l, int r) {
    o = ++nn;
    // cout<<o<<" "<<l<<" "<<r<<endl;
    if (l == r) {
      tn[o].ls = tn[o].rs = tn[o].cnt = tn[o].sum = 0;
      return;
    }
    int mid = (l + r) >> 1;
    build(tn[o].ls, l, mid);
    build(tn[o].rs, mid + 1, r);
  }
  TN query(int o, int l, int r) {
    if (l == r)
      return tn[o];
    int mid = (l + r) >> 1;
    if (P <= mid)
      return query(tn[o].ls, l, mid);
    else
      return query(tn[o].rs, mid + 1, r);
  }
  int add(int &o, int l, int r) {
    int no = ++nn;
    tn[no] = tn[o];
    o = no;
    tn[o].cnt++;
    tn[o].sum += V;
    if (l == r)
      return o;
    int mid = (l + r) >> 1;
    if (P <= mid)
      add(tn[o].ls, l, mid);
    else
      add(tn[o].rs, mid + 1, r);
    return o;
  }
  void show(int u) {
    cout << endl << "show tree " << u << endl;
    for (int i = 0; i <= n; i++) {
      P = i;
      TN node = query(u, 0, n);
      cout << i << " " << node.cnt << " " << node.sum << endl;
    }
    cout << endl;
  }
} psegt;

int vis[M];
int sum[M];
int root[M];

struct Edge {
  int v, c, d;
  Edge(int v = 0, int c = 0, int d = 0) : v(v), c(c), d(d) {}
};
vector<Edge> g[M];

void dfs(int u, int fa, int c, int d, int nsum) {
  sum[u] = nsum;
  psegt.P = c;
  psegt.V = d;
  root[u] = root[fa];
  psegt.add(root[u], 0, n);
  for (auto e : g[u]) {
    int v = e.v;
    if (v == fa)
      continue;
    dfs(v, u, e.c, e.d, nsum + e.d);
  }
}

struct LCA {
  int fa[M][MB], dep[M];
  void dfs_fa(int u, int f, int depth) {
    // cout<<u<<endl;
    dep[u] = depth;
    for (Edge e : g[u]) {
      int v = e.v;
      if (v == f)
        continue;
      if (!fa[v][0]) {
        fa[v][0] = u;
        dfs_fa(v, u, depth + 1);
      }
    }
  }
  void init_lca(int root) {
    memset(fa, 0, sizeof(fa));
    dfs_fa(root, 0, 0);
    // cout<<"here"<<endl;
    for (int j = 1; j < MB; j++) {
      for (int i = 1; i <= n; i++) {
        // cout<<i<<" "<<fa[i][j-1]<<" "<<fa[fa[i][j-1]][j-1]<<endl;
        fa[i][j] = fa[fa[i][j - 1]][j - 1];
      }
      // cout<<j<<endl;
    }
  }
  int lca(int a, int b) {
    if (dep[a] > dep[b])
      swap(a, b);
    int k = dep[b] - dep[a];
    for (int i = 0; (1 << i) <= k; i++)
      if (k & (1 << i))
        b = fa[b][i];
    if (a == b)
      return a;
    for (int i = MB - 1; i >= 0; i--) {
      if (fa[a][i] != fa[b][i]) {
        a = fa[a][i];
        b = fa[b][i];
      }
    }
    return fa[a][0];
  }
} lca;

int getsum(int u, int c, int d) {
  psegt.P = c;
  // cout<<u<<" c = "<<c<<" d = "<<d<<endl;
  // psegt.show(root[u]);
  TN node = psegt.query(root[u], 0, n);
  int ans = sum[u] - node.sum + node.cnt * d;
  // cout<<u<<" "<<sum[u]<<" "<<node.cnt<<" "<<node.sum;
  // cout<<" "<<ans<<endl;
  return ans;
}

int main() {
  // freopen("F_in.txt","r",stdin);
  scanf("%d%d", &n, &q);
  int u, v, c, d;
  for (int i = 1; i < n; i++) {
    scanf("%d%d%d%d", &u, &v, &c, &d);
    g[u].emplace_back(v, c, d);
    g[v].emplace_back(u, c, d);
  }
  root[0] = 1;
  psegt.init(n);
  // cout<<"init ok1"<<endl;
  dfs(1, 0, 0, 0, 0);
  // cout<<"init ok2"<<endl;
  lca.init_lca(1);
  // cout<<"init ok3"<<endl;
  while (q--) {
    scanf("%d%d%d%d", &c, &d, &u, &v);
    int l = lca.lca(u, v);
    // cout<<l<<endl;
    printf("%d\n", getsum(u, c, d) + getsum(v, c, d) - 2 * getsum(l, c, d));
  }
  return 0;
}

Fixed submission:
#include <algorithm>
#include <iostream>
#include <stdio.h>
#include <string.h>
#include <vector>
using namespace std;

const int M = 1e5 + 5;
const int MB = 20;

int n, q;

struct TN {
  int ls, rs;
  int cnt, sum;
  TN(int ls = 0, int rs = 0, int cnt = 0, int sum = 0)
      : ls(ls), rs(rs), cnt(cnt), sum(sum) {}
};

struct PSegT {
  TN tn[M * MB];
  int nn; // number of node
  int P, V;
  void init(int n) {
    nn = 0;
    int root;
    build(root, 0, n);
  }
  void build(int &o, int l, int r) {
    o = ++nn;
    // cout<<o<<" "<<l<<" "<<r<<endl;
    if (l == r) {
      tn[o].ls = tn[o].rs = tn[o].cnt = tn[o].sum = 0;
      return;
    }
    int mid = (l + r) >> 1;
    build(tn[o].ls, l, mid);
    build(tn[o].rs, mid + 1, r);
  }
  TN query(int o, int l, int r) {
    if (l == r)
      return tn[o];
    int mid = (l + r) >> 1;
    if (P <= mid)
      return query(tn[o].ls, l, mid);
    else
      return query(tn[o].rs, mid + 1, r);
  }
  int add(int &o, int l, int r) {
    int no = ++nn;
    tn[no] = tn[o];
    o = no;
    tn[o].cnt++;
    tn[o].sum += V;
    if (l == r)
      return o;
    int mid = (l + r) >> 1;
    if (P <= mid)
      add(tn[o].ls, l, mid);
    else
      add(tn[o].rs, mid + 1, r);
    return o;
  }
  void show(int u) {
    cout << endl << "show tree " << u << endl;
    for (int i = 0; i <= n; i++) {
      P = i;
      TN node = query(u, 0, n);
      cout << i << " " << node.cnt << " " << node.sum << endl;
    }
    cout << endl;
  }
} psegt;

int vis[M];
int sum[M];
int root[M];

struct Edge {
  int v, c, d;
  Edge(int v = 0, int c = 0, int d = 0) : v(v), c(c), d(d) {}
};
vector<Edge> g[M];

void dfs(int u, int fa, int c, int d, int nsum) {
  sum[u] = nsum;
  psegt.P = c;
  psegt.V = d;
  root[u] = root[fa];
  psegt.add(root[u], 0, n);
  for (auto e : g[u]) {
    int v = e.v;
    if (v == fa)
      continue;
    dfs(v, u, e.c, e.d, nsum + e.d);
  }
}

struct LCA {
  int fa[M][MB], dep[M];
  void dfs_fa(int u, int f, int depth) {
    // cout<<u<<endl;
    dep[u] = depth;
    for (Edge e : g[u]) {
      int v = e.v;
      if (v == f)
        continue;
      if (!fa[v][0]) {
        fa[v][0] = u;
        dfs_fa(v, u, depth + 1);
      }
    }
  }
  void init_lca(int root) {
    memset(fa, 0, sizeof(fa));
    dfs_fa(root, 0, 0);
    // cout<<"here"<<endl;
    for (int j = 1; j < MB; j++) {
      for (int i = 1; i <= n; i++) {
        // cout<<i<<" "<<fa[i][j-1]<<" "<<fa[fa[i][j-1]][j-1]<<endl;
        fa[i][j] = fa[fa[i][j - 1]][j - 1];
      }
      // cout<<j<<endl;
    }
  }
  int lca(int a, int b) {
    if (dep[a] > dep[b])
      swap(a, b);
    int k = dep[b] - dep[a];
    for (int i = 0; (1 << i) <= k; i++)
      if (k & (1 << i))
        b = fa[b][i];
    if (a == b)
      return a;
    for (int i = MB - 1; i >= 0; i--) {
      if (fa[a][i] != fa[b][i]) {
        a = fa[a][i];
        b = fa[b][i];
      }
    }
    return fa[a][0];
  }
} lca;

int getsum(int u, int c, int d) {
  psegt.P = c;
  // cout<<u<<" c = "<<c<<" d = "<<d<<endl;
  // psegt.show(root[u]);
  TN node = psegt.query(root[u], 0, n);
  int ans = sum[u] - node.sum + node.cnt * d;
  // cout<<u<<" "<<sum[u]<<" "<<node.cnt<<" "<<node.sum;
  // cout<<" "<<ans<<endl;
  return ans;
}

int main() {
  // freopen("F_in.txt","r",stdin);
  scanf("%d%d", &n, &q);
  int u, v, c, d;
  for (int i = 1; i < n; i++) {
    scanf("%d%d%d%d", &u, &v, &c, &d);
    g[u].emplace_back(v, c, d);
    g[v].emplace_back(u, c, d);
  }
  root[0] = 1;
  psegt.init(n);
  // cout<<"init ok1"<<endl;
  dfs(1, 0, 0, 0, 0);
  // cout<<"init ok2"<<endl;
  lca.init_lca(1);
  // cout<<"init ok3"<<endl;
  while (q--) {
    scanf("%d%d%d%d", &c, &d, &u, &v);
    int l = lca.lca(u, v);
    // cout<<l<<endl;
    printf("%d\n", getsum(u, c, d) + getsum(v, c, d) - 2 * getsum(l, c, d));
  }
  return 0;
}