Problem p02986 (C++)

Buggy submission (status: Time Limit Exceeded):
#pragma optimization_level 3
#pragma GCC optimize("Ofast,no-stack-protector,unroll-loops,fast-math")
#include <bits/stdc++.h>
#define F first
#define S second
#define vec vector
#define ms multiset
#define pb push_back
#define pll pair<ll, ll>
#define pdd pair<ld, ld>
#define pq priority_queue
#define umap unordered_map
#define uset unordered_set
#define pii pair<int, int>
#define pnn pair<Node *, Node *>
#define uid uniform_int_distribution
#define FILE                                                                   \
  ifstream in("board.in");                                                     \
  ofstream out("board.out");
#define fast                                                                   \
  cin.tie(0);                                                                  \
  cout.tie(0);                                                                 \
  cin.sync_with_stdio(0);                                                      \
  cout.sync_with_stdio(0);
using namespace std;
typedef string str;
typedef long long ll;
typedef long double ld;
mt19937 rnd(chrono::steady_clock::now()
                .time_since_epoch()
                .count()); // uid<int> u1(5, 10); u1(rnd);

struct edg {
  ll y, c, d;

  edg(int a, int b, int cc) {
    y = a;
    c = b;
    d = cc;
  }
};

struct Node {
  Node *l = 0, *r = 0;
  pii p = {0, 0};

  Node(Node *a, Node *b, pii c) {
    l = a;
    r = b;
    p = c;
  }

  Node(Node *a, Node *b) {
    l = a;
    r = b;
    p = {a->p.F + b->p.F, a->p.S + b->p.S};
  }

  Node(Node *n) {
    l = n->l;
    r = n->r;
    p = n->p;
  }
};

Node *add(ll l, ll r, ll pos, Node *n, ll val) {
  Node *cp = n ? new Node(n) : new Node(0, 0, {0, 0});
  if (l == r) {
    cp->p.F++;
    cp->p.S += val;
    return cp;
  }
  ll md = (l + r) / 2;
  if (pos <= md)
    cp->l = add(l, md, pos, cp->l, val);
  else
    cp->r = add(md + 1, r, pos, cp->r, val);
  return cp;
}

pll get(ll l, ll r, ll pos, Node *n) {
  if (!n)
    return {0, 0};
  if (l == r)
    return n->p;
  int md = (l + r) / 2;
  if (pos <= md)
    return get(l, md, pos, n->l);
  else
    return get(md + 1, r, pos, n->r);
}

ll a, tim = 0, ndo = -1, SZ = 1e5 + 5;
vec<vec<edg>> l;
vec<ll> sz, dep, tin, tout, numdo, posindo, eul, p, lf, rg, cl, dl, prs,
    firindo;
vec<Node *> m;

ll init1(ll v, ll pr, ll d, ll color, ll dln) {
  sz[v] = 1;
  if (pr != -1)
    cl[v] = color, dl[v] = dln, prs[v] = prs[pr] + dl[v];
  p[v] = pr;
  dep[v] = d;
  for (edg h : l[v]) {
    if (h.y != pr)
      sz[v] += init1(h.y, v, d + 1, h.c, h.d);
  }
  sort(l[v].begin(), l[v].end(), [](edg i1, edg i2) {
    return sz[i1.y] != sz[i2.y] ? sz[i1.y] > sz[i2.y] : i1.y < i2.y;
  });
  return sz[v];
}

void init2(ll v, int pr, int pdo) {
  eul.pb(v);
  if (pdo == 0)
    ndo++, lf.pb(eul.size() - 1);
  numdo[v] = ndo;
  tin[v] = tim++;
  posindo[v] = pdo;
  int cnt = 0;
  for (edg h : l[v]) {
    if (h.y == pr)
      continue;
    init2(h.y, v, (cnt == 0 ? pdo + 1 : 0));
    cnt++;
  }
  if (rg.size() < lf.size())
    rg.pb(eul.size() - 1);
  tout[v] = tim - 1;
}

ll check(ll v, ll pt) { return tin[v] <= tin[pt] && tout[pt] <= tout[v]; }

ll sm = 0;
pll go(ll x, ll y, ll c) {
  ll cnt = 0, d = 0, sx = x, sy = y;
  for (; !check(firindo[x], y);) {
    pll pp = get(0, SZ, c, m[lf[numdo[x]] + posindo[x]]);
    cnt += pp.F, d += pp.S;
    x = p[firindo[x]];
  }
  for (; numdo[x] != numdo[y];) {
    pll pp = get(0, SZ, c, m[lf[numdo[y]] + posindo[y]]);
    cnt += pp.F, d += pp.S;
    y = p[firindo[y]];
  }
  if (dep[x] > dep[y])
    swap(x, y);
  sm = prs[sx] + prs[sy] - 2 * prs[x];
  ll DO = numdo[x];
  pll p1 = get(0, SZ, c, m[lf[DO] + posindo[y]]);
  pll p2 = get(0, SZ, c, m[lf[DO] + posindo[x]]);
  return {cnt + p1.F - p2.F, d + p1.S - p2.S};
}

ll ga(ll x, ll y, ll c, ll dd) {
  ll summ = 0, cnt = 0, ss = 0;
  for (; !check(x, y);) {
    summ += dl[x];
    if (cl[x] == c)
      cnt++, ss += dl[x];
    x = p[x];
  }
  for (; y != x;) {
    summ += dl[y];
    if (cl[y] == c)
      cnt++, ss += dl[y];
    y = p[y];
  }
  return summ - ss + cnt * dd;
}

int main() {
  fast;
  ll z;
  cin >> a >> z;
  p = vec<ll>(a);
  sz = vec<ll>(a);
  cl = vec<ll>(a);
  dl = vec<ll>(a);
  dep = vec<ll>(a);
  tin = vec<ll>(a);
  prs = vec<ll>(a);
  tout = vec<ll>(a);
  numdo = vec<ll>(a);
  posindo = vec<ll>(a);
  firindo = vec<ll>(a);
  l = vec<vec<edg>>(a, vec<edg>());
  for (ll q = 0; q < a - 1; q++) {
    ll x, y, c, d;
    cin >> x >> y >> c >> d;
    x--;
    y--;
    if (x > y)
      swap(x, y);
    l[x].pb(edg(y, c, d));
    l[y].pb(edg(x, c, d));
  }
  init1(0, -1, 0, 0, 0);
  init2(0, -1, 0);
  ndo++;
  m = vec<Node *>(a);
  for (ll q = 0; q < ndo; q++) {
    m[lf[q]] = new Node(0, 0, {0, 0});
    for (ll w = lf[q]; w <= rg[q]; w++)
      m[w] = add(0, SZ, cl[eul[w]], (w == lf[q] ? m[w] : m[w - 1]), dl[eul[w]]),
      firindo[eul[w]] = eul[lf[q]];
  }
  // for(int q=0; q<ndo; q++) cout<<lf[q]<<" "<<rg[q]<<endl;
  for (; z; z--) {
    ll x, y, c, d;
    cin >> c >> d >> x >> y;
    x--;
    y--;
    sm = 0;
    // pll p = go(x,y,c);
    cout << ga(x, y, c, d) << "\n";
    // cout<<sm-p.S+p.F*d<<endl;
  }
}

Fixed submission:
#pragma optimization_level 3
#pragma GCC optimize("Ofast,no-stack-protector,unroll-loops,fast-math")
#include <bits/stdc++.h>
#define F first
#define S second
#define vec vector
#define ms multiset
#define pb push_back
#define pll pair<ll, ll>
#define pdd pair<ld, ld>
#define pq priority_queue
#define umap unordered_map
#define uset unordered_set
#define pii pair<int, int>
#define pnn pair<Node *, Node *>
#define uid uniform_int_distribution
#define FILE                                                                   \
  ifstream in("board.in");                                                     \
  ofstream out("board.out");
#define fast                                                                   \
  cin.tie(0);                                                                  \
  cout.tie(0);                                                                 \
  cin.sync_with_stdio(0);                                                      \
  cout.sync_with_stdio(0);
using namespace std;
typedef string str;
typedef long long ll;
typedef long double ld;
mt19937 rnd(chrono::steady_clock::now()
                .time_since_epoch()
                .count()); // uid<int> u1(5, 10); u1(rnd);

struct edg {
  ll y, c, d;

  edg(int a, int b, int cc) {
    y = a;
    c = b;
    d = cc;
  }
};

struct Node {
  Node *l = 0, *r = 0;
  pii p = {0, 0};

  Node(Node *a, Node *b, pii c) {
    l = a;
    r = b;
    p = c;
  }

  Node(Node *a, Node *b) {
    l = a;
    r = b;
    p = {a->p.F + b->p.F, a->p.S + b->p.S};
  }

  Node(Node *n) {
    l = n->l;
    r = n->r;
    p = n->p;
  }
};

Node *add(ll l, ll r, ll pos, Node *n, ll val) {
  Node *cp = n ? new Node(n) : new Node(0, 0, {0, 0});
  if (l == r) {
    cp->p.F++;
    cp->p.S += val;
    return cp;
  }
  ll md = (l + r) / 2;
  if (pos <= md)
    cp->l = add(l, md, pos, cp->l, val);
  else
    cp->r = add(md + 1, r, pos, cp->r, val);
  return cp;
}

pll get(ll l, ll r, ll pos, Node *n) {
  if (!n)
    return {0, 0};
  if (l == r)
    return n->p;
  int md = (l + r) / 2;
  if (pos <= md)
    return get(l, md, pos, n->l);
  else
    return get(md + 1, r, pos, n->r);
}

ll a, tim = 0, ndo = -1, SZ = 1e5 + 5;
vec<vec<edg>> l;
vec<ll> sz, dep, tin, tout, numdo, posindo, eul, p, lf, rg, cl, dl, prs,
    firindo;
vec<Node *> m;

ll init1(ll v, ll pr, ll d, ll color, ll dln) {
  sz[v] = 1;
  if (pr != -1)
    cl[v] = color, dl[v] = dln, prs[v] = prs[pr] + dl[v];
  p[v] = pr;
  dep[v] = d;
  for (edg h : l[v]) {
    if (h.y != pr)
      sz[v] += init1(h.y, v, d + 1, h.c, h.d);
  }
  sort(l[v].begin(), l[v].end(), [](edg i1, edg i2) {
    return sz[i1.y] != sz[i2.y] ? sz[i1.y] > sz[i2.y] : i1.y < i2.y;
  });
  return sz[v];
}

void init2(ll v, int pr, int pdo) {
  eul.pb(v);
  if (pdo == 0)
    ndo++, lf.pb(eul.size() - 1);
  numdo[v] = ndo;
  tin[v] = tim++;
  posindo[v] = pdo;
  int cnt = 0;
  for (edg h : l[v]) {
    if (h.y == pr)
      continue;
    init2(h.y, v, (cnt == 0 ? pdo + 1 : 0));
    cnt++;
  }
  if (rg.size() < lf.size())
    rg.pb(eul.size() - 1);
  tout[v] = tim - 1;
}

ll check(ll v, ll pt) { return tin[v] <= tin[pt] && tout[pt] <= tout[v]; }

ll sm = 0;
pll go(ll x, ll y, ll c) {
  ll cnt = 0, d = 0, sx = x, sy = y;
  for (; !check(firindo[x], y);) {
    pll pp = get(0, SZ, c, m[lf[numdo[x]] + posindo[x]]);
    cnt += pp.F, d += pp.S;
    x = p[firindo[x]];
  }
  for (; numdo[x] != numdo[y];) {
    pll pp = get(0, SZ, c, m[lf[numdo[y]] + posindo[y]]);
    cnt += pp.F, d += pp.S;
    y = p[firindo[y]];
  }
  if (dep[x] > dep[y])
    swap(x, y);
  sm = prs[sx] + prs[sy] - 2 * prs[x];
  ll DO = numdo[x];
  pll p1 = get(0, SZ, c, m[lf[DO] + posindo[y]]);
  pll p2 = get(0, SZ, c, m[lf[DO] + posindo[x]]);
  return {cnt + p1.F - p2.F, d + p1.S - p2.S};
}

ll ga(ll x, ll y, ll c, ll dd) {
  ll summ = 0, cnt = 0, ss = 0;
  for (; !check(x, y);) {
    summ += dl[x];
    if (cl[x] == c)
      cnt++, ss += dl[x];
    x = p[x];
  }
  for (; y != x;) {
    summ += dl[y];
    if (cl[y] == c)
      cnt++, ss += dl[y];
    y = p[y];
  }
  return summ - ss + cnt * dd;
}

int main() {
  fast;
  ll z;
  cin >> a >> z;
  p = vec<ll>(a);
  sz = vec<ll>(a);
  cl = vec<ll>(a);
  dl = vec<ll>(a);
  dep = vec<ll>(a);
  tin = vec<ll>(a);
  prs = vec<ll>(a);
  tout = vec<ll>(a);
  numdo = vec<ll>(a);
  posindo = vec<ll>(a);
  firindo = vec<ll>(a);
  l = vec<vec<edg>>(a, vec<edg>());
  for (ll q = 0; q < a - 1; q++) {
    ll x, y, c, d;
    cin >> x >> y >> c >> d;
    x--;
    y--;
    if (x > y)
      swap(x, y);
    l[x].pb(edg(y, c, d));
    l[y].pb(edg(x, c, d));
  }
  init1(0, -1, 0, 0, 0);
  init2(0, -1, 0);
  ndo++;
  m = vec<Node *>(a);
  for (ll q = 0; q < ndo; q++) {
    m[lf[q]] = new Node(0, 0, {0, 0});
    for (ll w = lf[q]; w <= rg[q]; w++)
      m[w] = add(0, SZ, cl[eul[w]], (w == lf[q] ? m[w] : m[w - 1]), dl[eul[w]]),
      firindo[eul[w]] = eul[lf[q]];
  }
  // for(int q=0; q<ndo; q++) cout<<lf[q]<<" "<<rg[q]<<endl;
  for (; z; z--) {
    ll x, y, c, d;
    cin >> c >> d >> x >> y;
    x--;
    y--;
    sm = 0;
    pll p = go(x, y, c);
    // cout<<ga(x,y,c,d)<<"\n";
    cout << sm - p.S + p.F * d << "\n";
  }
}