Problem p01771 (C++) — buggy submission, verdict: Time Limit Exceeded

// ?????¬???????????¬??????

#include <algorithm>
#include <cassert>
#include <cfloat>
#include <climits>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <ctime>
#include <deque>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>
using namespace std;

#define rep(i, a, n) for (int(i) = (a); (i) < (n); (i)++)
#define repq(i, a, n) for (int(i) = (a); (i) <= (n); (i)++)
#define repr(i, a, n) for (int(i) = (a); (i) >= (n); (i)--)
#define int long long int

template <typename T> void chmax(T &a, T b) { a = max(a, b); }
template <typename T> void chmin(T &a, T b) { a = min(a, b); }
template <typename T> void chadd(T &a, T b) { a = a + b; }

typedef pair<int, int> pii;
typedef long long ll;

int dx[] = {0, 0, 1, -1};
int dy[] = {1, -1, 0, 0};
constexpr ll INF = 1001001001001001LL;
constexpr ll MOD = 1000000007LL;

// ?????°??¨??¨?????¢??°

int upd(int a, int b, int l, int r) { return a + (r - l) * b; }

int qry(int a, int b) { return a + b; }

// Abstract Lazy Segment Tree (????±?????±?????????¶?????°???????????¨)
// Verified: AOJ ????????°??¨ verify ?????¨??¨

template <typename T> struct lazysegtree {
  // ?????????????????????
  vector<T> node, lazy;
  vector<int> need_upd;
  T I;
  int SIZE, M;

  // ????????¬????????§??? (update, query ??? 2 ???????????????)
  // update function ??????????????????????????¢????????????
  // upd_f(X, Y, l, r) -> ????????? [l, r) ??§???????????????????????? X ??? Y
  // ???????????? lazy ???????????? update ????????¨??????????????? 1
  // ??????????????¨?????°?????????
  T (*upd_f)(T, T, int, int), (*qry_f)(T, T);

  // ???????????¨?????????????????????????????¨?????? node ??¨ lazy
  // ??????????????§?????????
  lazysegtree(T (*op1)(T, T, int, int), T (*op2)(T, T), T X, int sz) {
    upd_f = op1;
    qry_f = op2;
    I = X;
    M = sz;
    SIZE = 1;
    while (SIZE < sz)
      SIZE *= 2;
    node.assign(2 * SIZE, I);
    lazy.assign(2 * SIZE, I);
    need_upd.assign(2 * SIZE, false);
  }

  int size() { return M; }

  // ?????? vec ????????§?????????
  void init(vector<T> vec) {
    int N = (int)vec.size();
    for (int i = 0; i < N; i++) {
      node[SIZE - 1 + i] = vec[i];
    }
    for (int i = SIZE - 2; i >= 0; i--) {
      node[i] = qry_f(node[2 * i + 1], node[2 * i + 2]);
    }
  }

  void lazy_eval(int k, int l, int r) {
    if (!need_upd[k])
      return;
    node[k] = upd_f(node[k], lazy[k], l, r);
    if (r - l > 1) {
      lazy[2 * k + 1] = upd_f(lazy[2 * k + 1], lazy[k], 0, 1);
      lazy[2 * k + 2] = upd_f(lazy[2 * k + 2], lazy[k], 0, 1);
      need_upd[2 * k + 1] = need_upd[2 * k + 2] = true;
    }
    lazy[k] = I;
    need_upd[k] = false;
  }

  // ???????????? [a, b) ??????????????? val ??????????????????
  // (upd_f ?????¨????????????)
  void update(int a, int b, T val, int l = 0, int r = -1, int k = 0) {
    if (r < 0)
      r = SIZE;
    lazy_eval(k, l, r);
    if (b <= l || r <= a)
      return;
    if (a <= l && r <= b) {
      lazy[k] = upd_f(lazy[k], val, 0, 1);
      need_upd[k] = true;
      lazy_eval(k, l, r);
    } else {
      int mid = (l + r) / 2;
      update(a, b, val, l, mid, 2 * k + 1);
      update(a, b, val, mid, r, 2 * k + 2);
      node[k] = qry_f(node[2 * k + 1], node[2 * k + 2]);
    }
  }

  // ???????????? [a, b) ?????????????????¨???????????????
  // (qry_f ?????¨????????????)
  T query(int a, int b, int l = 0, int r = -1, int k = 0) {
    if (r < 0)
      r = SIZE;
    lazy_eval(k, l, r);
    if (b <= l || r <= a)
      return I;
    if (a <= l && r <= b)
      return node[k];
    int mid = (l + r) / 2;
    T vl = query(a, b, l, mid, 2 * k + 1);
    T vr = query(a, b, mid, r, 2 * k + 2);
    return qry_f(vl, vr);
  }
};

// HL ????§£
// ???????????? v ???????????? c_i
// ?????????????????????????????¨????????¨?????¨???????????°???
// ?????§???????????? 1 ?????? c_h ??¨?????????????????¨????????? (v, c_h) ???
// "heavy" ????????§????????¨?????????
// ????????\\???????????? "light" ????????§????????¨?????????

// ?????????????????? "heavy" ??¨ "light" ?????????????????¨?????¨??? "heavy"
// ??????????????????
// (????????????????????¨ chain) ???????§£ ("heavy"
// ??§???????????£???????????????????????¨???????????¨??????) ??§?????????
// ????????? HL ????§£??¨?????????

// chain ???????´????????????¨?????¨???????????? O(log N) ??§?????????
// ????????????"heavy" ?????????????????????????????? v ??¨ "light"
// ????????§?????°???????????? v ?????????????????????????????????
// ??¨?????¨?????§????????? subsize(v) / 2
// ??\\?????¨?????£??????????????¨???????°?????????? "light"
// ??????????????????????????¨?????¨???????????????????????\\???????????????????????±??????
// O(log N) ??¨?????????

// ?§?????????¨????????¨???????????????????¨???????????§???????
template <typename T> struct Edge {
  int from, to;
  T cost;
  Edge(int s, T d) : to(s), cost(d) {}
  Edge(int f, int s, T d) : from(f), to(s), cost(d) {}

  bool operator<(const Edge &e) const { return cost < e.cost; }
  bool operator>(const Edge &e) const { return cost > e.cost; }
};

template <typename T> using Graph = vector<vector<Edge<T>>>;

struct HLD {
  int N;
  const Graph<int> G;
  // ???????????¨????????????
  // ??????????????±????????????????????????????????????????????¨????????¨?????¨????????????
  // ???????????????????????????????????? "heavy" ??????
  vector<int> depth, parent, subsize;
  // ???chain ????????????
  // chain ?????????????´????????°?????´??????????????????? 1
  // ????¬??????????????´????chain ?????§?????????????????????
  vector<int> head, last, prev, next, chain, idx;
  // chain ????????± (?????? vector ????????????????´??????????????????? chain
  // ????±???????)
  vector<vector<int>> chains;

  // ??? chain ????????????????????°??¨??? vector
  vector<lazysegtree<int>> segs;

  HLD(const Graph<int> &H, int r = -1)
      : N(H.size()), G(H), depth(N, -1), parent(N, 0), subsize(N, 0), head(N),
        last(N), prev(N, -1), next(N, -1), chain(N, -1), idx(N, 0) {
    if (r != -1)
      decompose(r);
  }

  // root ????????¨??????????§£
  void decompose(const int root) {
    stack<int> st;
    st.push(root);
    parent[root] = -1;
    depth[root] = 0;

    while (st.size()) {
      int cur_v = st.top();
      st.pop();
      // ?????????????????????????¨???????
      if (cur_v >= 0) {
        st.push(~cur_v);
        for (auto e : G[cur_v]) {
          // ??????????????§???????????°????????´??°
          if (depth[e.to] != -1)
            continue;
          depth[e.to] = depth[cur_v] + 1;
          parent[e.to] = cur_v;
          st.push(e.to);
        }
      }
      // ??°?????????
      else {
        int ma = 0;
        cur_v = ~cur_v;
        subsize[cur_v] = 1;
        for (auto e : G[cur_v]) {
          if (parent[cur_v] == e.to)
            continue;
          subsize[cur_v] += subsize[e.to];
          if (ma < subsize[e.to]) {
            // cur_v ??¨??¨?????¨????????????????????§??????????????¶
            // (????????? "heavy" ??????)
            ma = subsize[e.to];
            next[cur_v] = e.to;
          }
        }
      }
    }

    st.push(root);
    while (st.size()) {
      int cur_v = st.top();
      st.pop();
      for (auto e : G[cur_v]) {
        if (parent[cur_v] != e.to)
          st.push(e.to);
      }

      if (chain[cur_v] != -1)
        continue;
      chains.push_back(vector<int>());
      vector<int> &path = chains.back();

      for (int v = cur_v; v != -1; v = next[v])
        path.push_back(v);

      for (size_t i = 0; i < path.size(); i++) {
        int v = path[i];
        head[v] = path.front(), last[v] = path.back();
        prev[v] = (i != 0 ? path[i - 1] : -1);
        next[v] = (i + 1 != path.size() ? path[i + 1] : -1);
        chain[v] = (int)chains.size() - 1;
        idx[v] = i;
      }
    }
  }

  void buildSegtree() {
    segs.clear();
    for (size_t i = 0; i < chains.size(); i++) {
      segs.push_back(lazysegtree<int>(upd, qry, 0, chains[i].size()));
    }
  }

  // ?????? v ????????????????????? chain ????±?????????????
  // ???????????? chain ????????§????????????????????????????????? pair
  // ??§??????
  pair<int, int> get_index(int v) { return make_pair(chain[v], idx[v]); }

  // v ????±??????? chain ?????????????????? (head[v] ???????????? -1 ?????????)
  int climb(int v) { return parent[head[v]]; }

  // ?????? u ??¨ v ???????°???±????\\????
  int lca(int u, int v) {
    while (chain[u] != chain[v]) {
      if (depth[head[u]] < depth[head[v]])
        v = climb(v);
      else
        u = climb(u);
    }
    return depth[u] < depth[v] ? u : v;
  }

  void update(int v, int x) {
    int chain_id, v_id;
    tie(chain_id, v_id) = get_index(v);

    segs[chain_id].update(v_id, chains[chain_id].size(), x);
  }

  int solve(int v, int root) {
    int t = v, ans = 0;
    while (chain[t] != chain[root]) {
      int nxt = climb(t);

      pii info_t = get_index(t);
      ans += segs[info_t.first].query(0, info_t.second);

      if (nxt != -1) {
        int depth_diff = depth[v] - depth[nxt];
        int chain_id, v_id;
        tie(chain_id, v_id) = get_index(nxt);
        ans += depth_diff * segs[chain_id].query(v_id, v_id + 1);
      }
      t = nxt;
    }

    int chain_id = chain[t], v_id = idx[t], root_id = idx[root];
    ans += segs[chain_id].query(root_id, v_id);
    return ans;
  }

  int query(int u, int v) {
    int l = lca(u, v);
    return solve(u, 0) + solve(v, 0) - 2 * solve(l, 0);
  }
};

signed main() {
  int N, Q;
  scanf("%lld%lld", &N, &Q);

  Graph<int> G(N);
  rep(i, 0, N - 1) {
    int a, b;
    scanf("%lld%lld", &a, &b);
    // a < b ???????¨????????????????
    G[a].push_back(Edge<int>(b, 1));
  }

  HLD hl(G, 0);
  hl.buildSegtree();
  for (int i = 0; i < Q; i++) {
    int q, a, b;
    scanf("%lld%lld%lld", &q, &a, &b);
    if (q == 0) {
      // calc distance
      printf("%lld\\n", hl.query(a, b));
    } else {
      // modify the weight
      hl.update(a, b);
    }
  }
  return 0;
}

Error: TLE

Fixed submission:

// ?????¬???????????¬??????

#include <algorithm>
#include <cassert>
#include <cfloat>
#include <climits>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <ctime>
#include <deque>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>
using namespace std;

#define rep(i, a, n) for (int(i) = (a); (i) < (n); (i)++)
#define repq(i, a, n) for (int(i) = (a); (i) <= (n); (i)++)
#define repr(i, a, n) for (int(i) = (a); (i) >= (n); (i)--)
#define int long long int

template <typename T> void chmax(T &a, T b) { a = max(a, b); }
template <typename T> void chmin(T &a, T b) { a = min(a, b); }
template <typename T> void chadd(T &a, T b) { a = a + b; }

typedef pair<int, int> pii;
typedef long long ll;

int dx[] = {0, 0, 1, -1};
int dy[] = {1, -1, 0, 0};
constexpr ll INF = 1001001001001001LL;
constexpr ll MOD = 1000000007LL;

// ?????°??¨??¨?????¢??°

int upd(int a, int b, int l, int r) { return a + (r - l) * b; }

int qry(int a, int b) { return a + b; }

// Abstract Lazy Segment Tree (????±?????±?????????¶?????°???????????¨)
// Verified: AOJ ????????°??¨ verify ?????¨??¨

template <typename T> struct lazysegtree {
  // ?????????????????????
  vector<T> node, lazy;
  vector<int> need_upd;
  T I;
  int SIZE, M;

  // ????????¬????????§??? (update, query ??? 2 ???????????????)
  // update function ??????????????????????????¢????????????
  // upd_f(X, Y, l, r) -> ????????? [l, r) ??§???????????????????????? X ??? Y
  // ???????????? lazy ???????????? update ????????¨??????????????? 1
  // ??????????????¨?????°?????????
  T (*upd_f)(T, T, int, int), (*qry_f)(T, T);

  // ???????????¨?????????????????????????????¨?????? node ??¨ lazy
  // ??????????????§?????????
  lazysegtree(T (*op1)(T, T, int, int), T (*op2)(T, T), T X, int sz) {
    upd_f = op1;
    qry_f = op2;
    I = X;
    M = sz;
    SIZE = 1;
    while (SIZE < sz)
      SIZE *= 2;
    node.assign(2 * SIZE, I);
    lazy.assign(2 * SIZE, I);
    need_upd.assign(2 * SIZE, false);
  }

  int size() { return M; }

  // ?????? vec ????????§?????????
  void init(vector<T> vec) {
    int N = (int)vec.size();
    for (int i = 0; i < N; i++) {
      node[SIZE - 1 + i] = vec[i];
    }
    for (int i = SIZE - 2; i >= 0; i--) {
      node[i] = qry_f(node[2 * i + 1], node[2 * i + 2]);
    }
  }

  void lazy_eval(int k, int l, int r) {
    if (!need_upd[k])
      return;
    node[k] = upd_f(node[k], lazy[k], l, r);
    if (r - l > 1) {
      lazy[2 * k + 1] = upd_f(lazy[2 * k + 1], lazy[k], 0, 1);
      lazy[2 * k + 2] = upd_f(lazy[2 * k + 2], lazy[k], 0, 1);
      need_upd[2 * k + 1] = need_upd[2 * k + 2] = true;
    }
    lazy[k] = I;
    need_upd[k] = false;
  }

  // ???????????? [a, b) ??????????????? val ??????????????????
  // (upd_f ?????¨????????????)
  void update(int a, int b, T val, int l = 0, int r = -1, int k = 0) {
    if (r < 0)
      r = SIZE;
    lazy_eval(k, l, r);
    if (b <= l || r <= a)
      return;
    if (a <= l && r <= b) {
      lazy[k] = upd_f(lazy[k], val, 0, 1);
      need_upd[k] = true;
      lazy_eval(k, l, r);
    } else {
      int mid = (l + r) / 2;
      update(a, b, val, l, mid, 2 * k + 1);
      update(a, b, val, mid, r, 2 * k + 2);
      node[k] = qry_f(node[2 * k + 1], node[2 * k + 2]);
    }
  }

  // ???????????? [a, b) ?????????????????¨???????????????
  // (qry_f ?????¨????????????)
  T query(int a, int b, int l = 0, int r = -1, int k = 0) {
    if (r < 0)
      r = SIZE;
    lazy_eval(k, l, r);
    if (b <= l || r <= a)
      return I;
    if (a <= l && r <= b)
      return node[k];
    int mid = (l + r) / 2;
    T vl = query(a, b, l, mid, 2 * k + 1);
    T vr = query(a, b, mid, r, 2 * k + 2);
    return qry_f(vl, vr);
  }
};

// HL ????§£
// ???????????? v ???????????? c_i
// ?????????????????????????????¨????????¨?????¨???????????°???
// ?????§???????????? 1 ?????? c_h ??¨?????????????????¨????????? (v, c_h) ???
// "heavy" ????????§????????¨?????????
// ????????\\???????????? "light" ????????§????????¨?????????

// ?????????????????? "heavy" ??¨ "light" ?????????????????¨?????¨??? "heavy"
// ??????????????????
// (????????????????????¨ chain) ???????§£ ("heavy"
// ??§???????????£???????????????????????¨???????????¨??????) ??§?????????
// ????????? HL ????§£??¨?????????

// chain ???????´????????????¨?????¨???????????? O(log N) ??§?????????
// ????????????"heavy" ?????????????????????????????? v ??¨ "light"
// ????????§?????°???????????? v ?????????????????????????????????
// ??¨?????¨?????§????????? subsize(v) / 2
// ??\\?????¨?????£??????????????¨???????°?????????? "light"
// ??????????????????????????¨?????¨???????????????????????\\???????????????????????±??????
// O(log N) ??¨?????????

// ?§?????????¨????????¨???????????????????¨???????????§???????
template <typename T> struct Edge {
  int from, to;
  T cost;
  Edge(int s, T d) : to(s), cost(d) {}
  Edge(int f, int s, T d) : from(f), to(s), cost(d) {}

  bool operator<(const Edge &e) const { return cost < e.cost; }
  bool operator>(const Edge &e) const { return cost > e.cost; }
};

template <typename T> using Graph = vector<vector<Edge<T>>>;

struct HLD {
  int N;
  const Graph<int> G;
  // ???????????¨????????????
  // ??????????????±????????????????????????????????????????????¨????????¨?????¨????????????
  // ???????????????????????????????????? "heavy" ??????
  vector<int> depth, parent, subsize;
  // ???chain ????????????
  // chain ?????????????´????????°?????´??????????????????? 1
  // ????¬??????????????´????chain ?????§?????????????????????
  vector<int> head, last, prev, next, chain, idx;
  // chain ????????± (?????? vector ????????????????´??????????????????? chain
  // ????±???????)
  vector<vector<int>> chains;

  // ??? chain ????????????????????°??¨??? vector
  vector<lazysegtree<int>> segs;

  HLD(const Graph<int> &H, int r = -1)
      : N(H.size()), G(H), depth(N, -1), parent(N, 0), subsize(N, 0), head(N),
        last(N), prev(N, -1), next(N, -1), chain(N, -1), idx(N, 0) {
    if (r != -1)
      decompose(r);
  }

  // root ????????¨??????????§£
  void decompose(const int root) {
    stack<int> st;
    st.push(root);
    parent[root] = -1;
    depth[root] = 0;

    while (st.size()) {
      int cur_v = st.top();
      st.pop();
      // ?????????????????????????¨???????
      if (cur_v >= 0) {
        st.push(~cur_v);
        for (auto e : G[cur_v]) {
          // ??????????????§???????????°????????´??°
          if (depth[e.to] != -1)
            continue;
          depth[e.to] = depth[cur_v] + 1;
          parent[e.to] = cur_v;
          st.push(e.to);
        }
      }
      // ??°?????????
      else {
        int ma = 0;
        cur_v = ~cur_v;
        subsize[cur_v] = 1;
        for (auto e : G[cur_v]) {
          if (parent[cur_v] == e.to)
            continue;
          subsize[cur_v] += subsize[e.to];
          if (ma < subsize[e.to]) {
            // cur_v ??¨??¨?????¨????????????????????§??????????????¶
            // (????????? "heavy" ??????)
            ma = subsize[e.to];
            next[cur_v] = e.to;
          }
        }
      }
    }

    st.push(root);
    while (st.size()) {
      int cur_v = st.top();
      st.pop();
      for (auto e : G[cur_v]) {
        if (parent[cur_v] != e.to)
          st.push(e.to);
      }

      if (chain[cur_v] != -1)
        continue;
      chains.push_back(vector<int>());
      vector<int> &path = chains.back();

      for (int v = cur_v; v != -1; v = next[v])
        path.push_back(v);

      for (size_t i = 0; i < path.size(); i++) {
        int v = path[i];
        head[v] = path.front(), last[v] = path.back();
        prev[v] = (i != 0 ? path[i - 1] : -1);
        next[v] = (i + 1 != path.size() ? path[i + 1] : -1);
        chain[v] = (int)chains.size() - 1;
        idx[v] = i;
      }
    }
  }

  void buildSegtree() {
    segs.clear();
    for (size_t i = 0; i < chains.size(); i++) {
      segs.push_back(lazysegtree<int>(upd, qry, 0, chains[i].size()));
    }
  }

  // ?????? v ????????????????????? chain ????±?????????????
  // ???????????? chain ????????§????????????????????????????????? pair
  // ??§??????
  pair<int, int> get_index(int v) { return make_pair(chain[v], idx[v]); }

  // v ????±??????? chain ?????????????????? (head[v] ???????????? -1 ?????????)
  int climb(int v) { return parent[head[v]]; }

  // ?????? u ??¨ v ???????°???±????\\????
  int lca(int u, int v) {
    while (chain[u] != chain[v]) {
      if (depth[head[u]] < depth[head[v]])
        v = climb(v);
      else
        u = climb(u);
    }
    return depth[u] < depth[v] ? u : v;
  }

  void update(int v, int x) {
    int chain_id, v_id;
    tie(chain_id, v_id) = get_index(v);

    segs[chain_id].update(v_id, chains[chain_id].size(), x);
  }

  int solve(int v, int root) {
    int t = v, ans = 0;
    while (chain[t] != chain[root]) {
      int nxt = climb(v);

      pii info_t = get_index(t);
      ans += segs[info_t.first].query(0, info_t.second);

      if (nxt != -1) {
        int depth_diff = depth[v] - depth[nxt];
        int chain_id, v_id;
        tie(chain_id, v_id) = get_index(nxt);
        ans += depth_diff * segs[chain_id].query(v_id, v_id + 1);
      }
      t = nxt;
    }

    int chain_id = chain[t], v_id = idx[t], root_id = idx[root];
    ans += segs[chain_id].query(root_id, v_id);
    return ans;
  }

  int query(int u, int v) {
    int l = lca(u, v);
    return solve(u, 0) + solve(v, 0) - 2 * solve(l, 0);
  }
};

signed main() {
  int N, Q;
  scanf("%lld%lld", &N, &Q);

  Graph<int> G(N);
  rep(i, 0, N - 1) {
    int a, b;
    scanf("%lld%lld", &a, &b);
    // a < b ???????¨????????????????
    G[a].push_back(Edge<int>(b, 1));
  }

  HLD hl(G, 0);
  hl.buildSegtree();
  while (Q--) {
    int q, a, b;
    scanf("%lld%lld%lld", &q, &a, &b);
    if (q == 0) {
      // calc distance
      printf("%lld\\n", hl.query(a, b));
    } else {
      // modify the weight
      hl.update(a, b);
    }
  }
  return 0;
}