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

#include <bits/stdc++.h>
using namespace std;

#define rep(i, x, y) for (int i = (x); i < (y); ++i)
#define debug(x) #x << "=" << (x)

#ifdef DEBUG
#define _GLIBCXX_DEBUG
#define show(x) std::cerr << debug(x) << " (L:" << __LINE__ << ")" << std::endl
#else
#define show(x)
#endif

typedef long long int ll;
typedef pair<int, int> pii;
template <typename T> using vec = std::vector<T>;

const int inf = 1 << 30;
const long long int infll = 1LL << 62;
const double eps = 1e-9;
const int dx[] = {1, 0, -1, 0}, dy[] = {0, 1, 0, -1};

template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  os << "[";
  for (const auto &v : vec) {
    os << v << ",";
  }
  os << "]";
  return os;
}

struct HLDecomposition {
  const vector<vector<int>> graph;
  const int root, n;
  vector<int> parent, subtree_size, depth, first, last, prev, next, path_id, at;
  vector<vector<int>> paths, children;

  HLDecomposition(const vector<vector<int>> &graph, int root)
      : graph(graph), root(root), n(graph.size()), parent(n),
        subtree_size(n, 1), depth(n), first(n), last(n), prev(n, -1),
        next(n, -1), path_id(n), at(n) {
    preprocess();
    decompose();
  }

  void preprocess() {
    stack<tuple<int, int, int>> dfs_stack;
    dfs_stack.push(make_tuple(root, -1, 0));
    while (!dfs_stack.empty()) {
      auto t = dfs_stack.top();
      dfs_stack.pop();
      int v = get<0>(t), p = get<1>(t), i = get<2>(t);
      if (i == 0) {
        if (v != root)
          depth[v] = depth[p] + 1;
        dfs_stack.push(make_tuple(v, p, 1));
        for (int u : graph[v]) {
          if (u == p)
            continue;
          dfs_stack.push(make_tuple(u, v, 0));
        }
      } else {
        parent[v] = p;
        for (int u : graph[v]) {
          if (u == p)
            continue;
          subtree_size[v] += subtree_size[u];
        }
      }
    }
  }

  void decompose() {
    int cnt = 0;
    stack<pii> dfs_stack;
    dfs_stack.push(make_pair(root, 0));
    while (!dfs_stack.empty()) {
      auto p = dfs_stack.top();
      dfs_stack.pop();

      int v = p.first, id = p.second;
      if (paths.size() <= cnt)
        paths.resize(cnt + 1);
      if (children.size() <= cnt)
        children.resize(cnt + 1);

      int cur = v;
      while (cur != -1) {
        paths[id].push_back(cur);
        int heavy_edge = -1;
        for (int u : graph[cur]) {
          if (u == parent[cur])
            continue;
          if (heavy_edge == -1 and subtree_size[u] * 2 >= subtree_size[cur] - 1)
            heavy_edge = u;
          else {
            ++cnt;
            children[id].push_back(cnt);
            dfs_stack.push(make_pair(u, cnt));
          }
        }
        cur = heavy_edge;
      }
      for (int i = 0; i < paths[id].size(); ++i) {
        int u = paths[id][i];
        first[u] = paths[id].front();
        last[u] = paths[id].back();
        if (i != 0)
          prev[u] = paths[id][i - 1];
        if (i != paths[id].size() - 1)
          next[u] = paths[id][i + 1];
        path_id[u] = id;
        at[u] = i;
      }
    }
  }

  int size() const { return n; }

  int go_up(int v) { return parent[first[v]]; }

  std::pair<int, int> get_index(int v) { return make_pair(path_id[v], at[v]); }

  int lca(int u, int v) {
    while (path_id[u] != path_id[v]) {
      if (depth[first[u]] < depth[first[v]])
        v = go_up(v);
      else
        u = go_up(u);
    }
    return depth[u] < depth[v] ? u : v;
  }

  int highest_child(int path_id, int idx) {
    int lb = -1, ub = children[path_id].size();
    while (ub - lb > 1) {
      int m = (lb + ub) / 2;
      if (at[go_up(paths[children[path_id][m]].front())] >= idx)
        ub = m;
      else
        lb = m;
    }
    return ub;
  };
};

// range add,get sum
template <class T> class segtree {
public:
  int n, size_;
  vector<T> dat, sum; //????????????,?????????????§?????¶?????????????
  segtree(int size_, T x) : size_(size_) {
    n = 1;
    while (n < size_)
      n *= 2;
    dat.assign(2 * n - 1, x);
    sum.assign(2 * n - 1, x);
  }
  int size() const { return size_; }
  void add(int a, int b, T x) { add(a, b, x, 0, 0, n); }
  T add(int a, int b, T x, int k, int l, int r) {
    if (b <= l or r <= a)
      return dat[k];
    if (a <= l and r <= b) {
      sum[k] += x;
      return dat[k] += x * (r - l);
    }
    int m = (l + r) / 2;
    return dat[k] = add(a, b, x, 2 * k + 1, l, m) +
                    add(a, b, x, 2 * k + 2, m, r) + sum[k] * (r - l);
  }
  T query(int a, int b) { return query(a, b, 0, 0, n); }
  T query(int a, int b, int k, int l, int r) {
    if (b <= l or r <= a)
      return 0;
    if (a <= l and r <= b)
      return dat[k];
    int m = (l + r) / 2;
    return query(a, b, 2 * k + 1, l, m) + query(a, b, 2 * k + 2, m, r) +
           sum[k] * (min(r, b) - max(l, a));
  }
};

void solve() {
  int n, q;
  cin >> n >> q;
  vector<vector<int>> graph(n);
  rep(i, 0, n - 1) {
    int a, b;
    cin >> a >> b;
    graph[a].push_back(b);
    graph[b].push_back(a);
  }

  HLDecomposition hld(graph, 0);
  vector<ll> all(hld.paths.size());
  vector<segtree<ll>> part, connections;
  rep(i, 0, hld.paths.size()) {
    part.push_back(segtree<ll>(hld.paths[i].size() - 1, 0));
    connections.push_back(segtree<ll>(hld.children[i].size(), 0));
  }
  rep(i, 0, q) {
    int t;
    cin >> t;
    if (!t) {
      int u, v, w;
      cin >> u >> v;
      w = hld.lca(u, v);
      function<ll(int, int)> rec = [&](int x, int d) {
        ll res = part[hld.path_id[x]].query(0, hld.at[x]);
        int y = hld.go_up(x);
        if (y == -1)
          return res;
        d += hld.at[x] + 1;
        int i = hld.highest_child(hld.path_id[y], hld.at[y]);
        return res + connections[hld.path_id[y]].query(i, i + 1) * d +
               rec(y, d);
      };
      cout << rec(u, 0) + rec(v, 0) - rec(w, 0) * 2 << endl;
    } else {
      int v;
      ll val;
      cin >> v >> val;
      int c = hld.path_id[v];
      part[c].add(hld.at[v], part[c].size(), val);
      connections[c].add(hld.highest_child(c, hld.at[v]), connections[c].size(),
                         val);
    }
  }
}

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  cout.setf(ios::fixed);
  cout.precision(10);
  solve();
  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

#define rep(i, x, y) for (int i = (x); i < (y); ++i)
#define debug(x) #x << "=" << (x)

#ifdef DEBUG
#define _GLIBCXX_DEBUG
#define show(x) std::cerr << debug(x) << " (L:" << __LINE__ << ")" << std::endl
#else
#define show(x)
#endif

typedef long long int ll;
typedef pair<int, int> pii;
template <typename T> using vec = std::vector<T>;

const int inf = 1 << 30;
const long long int infll = 1LL << 62;
const double eps = 1e-9;
const int dx[] = {1, 0, -1, 0}, dy[] = {0, 1, 0, -1};

template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  os << "[";
  for (const auto &v : vec) {
    os << v << ",";
  }
  os << "]";
  return os;
}

struct HLDecomposition {
  const vector<vector<int>> graph;
  const int root, n;
  vector<int> parent, subtree_size, depth, first, last, prev, next, path_id, at;
  vector<vector<int>> paths, children;

  HLDecomposition(const vector<vector<int>> &graph, int root)
      : graph(graph), root(root), n(graph.size()), parent(n),
        subtree_size(n, 1), depth(n), first(n), last(n), prev(n, -1),
        next(n, -1), path_id(n), at(n) {
    preprocess();
    show("ok1");
    decompose();
    show("ok2");
  }

  void preprocess() {
    stack<tuple<int, int, int>> dfs_stack;
    dfs_stack.push(make_tuple(root, -1, 0));
    while (!dfs_stack.empty()) {
      auto t = dfs_stack.top();
      dfs_stack.pop();
      int v = get<0>(t), p = get<1>(t), i = get<2>(t);
      if (i == 0) {
        if (v != root)
          depth[v] = depth[p] + 1;
        dfs_stack.push(make_tuple(v, p, 1));
        for (int u : graph[v]) {
          if (u == p)
            continue;
          dfs_stack.push(make_tuple(u, v, 0));
        }
      } else {
        parent[v] = p;
        for (int u : graph[v]) {
          if (u == p)
            continue;
          subtree_size[v] += subtree_size[u];
        }
      }
    }
  }

  void decompose() {
    int cnt = 0;
    stack<pii> dfs_stack;
    dfs_stack.push(make_pair(root, 0));
    while (!dfs_stack.empty()) {
      auto p = dfs_stack.top();
      dfs_stack.pop();

      int v = p.first, id = p.second;
      if (paths.size() <= cnt)
        paths.resize(cnt + 1);
      if (children.size() <= cnt)
        children.resize(cnt + 1);

      int cur = v;
      while (cur != -1) {
        paths[id].push_back(cur);
        int heavy_edge = -1;
        for (int u : graph[cur]) {
          if (u == parent[cur])
            continue;
          if (heavy_edge == -1 and subtree_size[u] * 2 >= subtree_size[cur] - 1)
            heavy_edge = u;
          else {
            ++cnt;
            children[id].push_back(cnt);
            dfs_stack.push(make_pair(u, cnt));
          }
        }
        cur = heavy_edge;
      }
      for (int i = 0; i < paths[id].size(); ++i) {
        int u = paths[id][i];
        first[u] = paths[id].front();
        last[u] = paths[id].back();
        if (i != 0)
          prev[u] = paths[id][i - 1];
        if (i != paths[id].size() - 1)
          next[u] = paths[id][i + 1];
        path_id[u] = id;
        at[u] = i;
      }
    }
  }

  int size() const { return n; }

  int go_up(int v) { return parent[first[v]]; }

  std::pair<int, int> get_index(int v) { return make_pair(path_id[v], at[v]); }

  int lca(int u, int v) {
    while (path_id[u] != path_id[v]) {
      if (depth[first[u]] < depth[first[v]])
        v = go_up(v);
      else
        u = go_up(u);
    }
    return depth[u] < depth[v] ? u : v;
  }
};

// range add,get sum
template <class T> class segtree {
public:
  int n, size_;
  vector<T> dat, sum; //????????????,?????????????§?????¶?????????????
  segtree(int size_, T x) : size_(size_) {
    n = 1;
    while (n < size_)
      n *= 2;
    dat.assign(2 * n - 1, x);
    sum.assign(2 * n - 1, x);
  }
  int size() const { return size_; }
  void add(int a, int b, T x) { add(a, b, x, 0, 0, n); }
  T add(int a, int b, T x, int k, int l, int r) {
    if (b <= l or r <= a)
      return dat[k];
    if (a <= l and r <= b) {
      sum[k] += x;
      return dat[k] += x * (r - l);
    }
    int m = (l + r) / 2;
    return dat[k] = add(a, b, x, 2 * k + 1, l, m) +
                    add(a, b, x, 2 * k + 2, m, r) + sum[k] * (r - l);
  }
  T query(int a, int b) { return query(a, b, 0, 0, n); }
  T query(int a, int b, int k, int l, int r) {
    if (b <= l or r <= a)
      return 0;
    if (a <= l and r <= b)
      return dat[k];
    int m = (l + r) / 2;
    return query(a, b, 2 * k + 1, l, m) + query(a, b, 2 * k + 2, m, r) +
           sum[k] * (min(r, b) - max(l, a));
  }
};

void solve() {
  int n, q;
  cin >> n >> q;
  vector<vector<int>> graph(n);
  rep(i, 0, n - 1) {
    int a, b;
    cin >> a >> b;
    graph[a].push_back(b);
    graph[b].push_back(a);
  }

  HLDecomposition hld(graph, 0);
  vector<segtree<ll>> vs;
  vector<map<int, ll>> cost(hld.paths.size());
  rep(i, 0, hld.paths.size())
      vs.push_back(segtree<ll>(hld.paths[i].size() - 1, 0));
  rep(i, 0, q) {
    int t;
    cin >> t;
    if (!t) {
      int u, v, w;
      cin >> u >> v;
      w = hld.lca(u, v);
      function<ll(int)> rec = [&](int x) {
        ll tmp = vs[hld.path_id[x]].query(0, hld.at[x]);
        if (hld.go_up(x) == -1)
          return tmp;
        return tmp + cost[hld.path_id[x]][hld.path_id[hld.go_up(x)]] +
               rec(hld.go_up(x));
      };
      cout << rec(u) + rec(v) - rec(w) * 2 << endl;
    } else {
      int v;
      ll val;
      cin >> v >> val;
      int c = hld.path_id[v];
      vs[c].add(hld.at[v], vs[c].size(), val);
      function<void(int)> add = [&](int x) {
        vs[x].add(0, vs[x].size(), val);
        for (int y : hld.children[x]) {
          cost[x][y] += val;
          cost[y][x] += val;
          add(y);
        }
      };
      for (int x : hld.children[c]) {
        if (hld.at[hld.go_up(hld.paths[x][0])] < hld.at[v])
          continue;
        cost[c][x] += val;
        cost[x][c] += val;
        add(x);
      }
    }
  }
}

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  cout.setf(ios::fixed);
  cout.precision(10);
  solve();
  return 0;
}