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

#include <algorithm>
#include <cstdio>
#include <iostream>
#include <queue>
#include <stack>
#include <vector>
using namespace std;
#define MAX_N 300005
#define MAX_M (1 << 20)
#define INF (1 << 28)

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

struct seg {
  ll t[MAX_M];
  ll u[MAX_M];
  void init() {
    for (int i = 0; i < MAX_M; i++)
      t[i] = u[i] = 0;
  }
  void add(int a, int b, ll value, int x, int l, int r) {
    if (b <= l || r <= a) {
      return;
    } else if (a <= l && r <= b) {
      t[x] += value;
      return;
    } else {
      add(a, b, value, x * 2 + 1, l, (l + r) / 2);
      add(a, b, value, x * 2 + 2, (l + r) / 2, r);
      u[x] += value * (min(r, b) - max(l, a));
    }
  }
  ll sum(int a, int b, int x, int l, int r) {
    if (b <= l || r <= a) {
      return 0;
    } else if (a <= l && r <= b) {
      return t[x] * (r - l) + u[x];
    } else {
      ll res = (min(r, b) - max(l, a)) * t[x];
      res += sum(a, b, x * 2 + 1, l, (l + r) / 2);
      res += sum(a, b, x * 2 + 2, (l + r) / 2, r);
      return res;
    }
  }
};

seg A, B;

int N, Tc;
int a, b, c;
vector<int> G[MAX_N];
int parent[MAX_N];
int depth[MAX_N];
int id[MAX_N];
int id2[MAX_N];
vector<int> g;
vector<int> al, ar;
vector<int> bl, br;
int ca = 0, cb = 0;
P dp[MAX_M][26];
stack<int> Q;
void lca_init() {
  for (int i = 0; i < MAX_N; i++)
    parent[i] = -1, id[i] = INF;

  depth[0] = 0;
  Q.push(0);

  while (!Q.empty()) {
    int pos = Q.top();
    Q.pop();
    id[pos] = g.size();
    g.push_back(pos);
    int size = G[pos].size();
    if (size && G[pos][size - 1] == parent[pos])
      G[pos].pop_back(), size--;
    if (size == 0)
      continue;
    Q.push(pos);
    Q.push(G[pos][size - 1]);
    depth[G[pos][size - 1]] = depth[pos] + 1;
    parent[G[pos][size - 1]] = pos;
    G[pos].pop_back();
  }

  /*
  for(int i=0;i<(int)g.size();i++){
    if(i){
      if(depth[g[i-1]]<depth[g[i]])cout<<'+';
      else cout<<'-';
    }
    cout<<g[i];
  }
  cout<<endl;
  */
  al.resize(g.size());
  ar.resize(g.size());
  bl.resize(g.size());
  br.resize(g.size());
  for (int i = 0; i < (int)g.size(); i++) {
    al[i] = bl[i] = INF;
    ar[i] = br[i] = -INF;
    id[g[i]] = min(id[g[i]], i);
    id2[g[i]] = i;
  }

  for (int i = 0; i + 1 < (int)g.size(); i++) {
    if (depth[g[i]] < depth[g[i + 1]]) {
      // cout<<' '<<ca;
      al[i] = ar[i] = ca;
      ca++;
    } else {
      // cout<<' '<<cb;
      bl[i] = br[i] = cb;
      cb++;
    }
  }
  // cout<<endl;

  for (int i = 1; i + 1 < (int)g.size(); i++) {
    ar[i] = max(ar[i], ar[i - 1]);
    br[i] = max(br[i], br[i - 1]);
  }
  for (int i = (int)g.size() - 2; i >= 0; i--) {
    al[i] = min(al[i], al[i + 1]);
    bl[i] = min(bl[i], bl[i + 1]);
  }

  // for(int i=0;i<(int)g.size();i++)cout<<' '<<al[i];cout<<endl;
  //  for(int i=0;i<(int)g.size();i++)cout<<' '<<ar[i];cout<<endl;
  // for(int i=0;i<(int)g.size();i++)cout<<' '<<bl[i];cout<<endl;
  // for(int i=0;i<(int)g.size();i++)cout<<' '<<br[i];;cout<<endl;

  for (int i = 0; i < (int)g.size(); i++)
    dp[i][0] = P(depth[g[i]], i);

  for (int k = 1, K = 1; k < 26; k++, K *= 2) {
    for (int i = 0; i < (int)g.size(); i++) {
      dp[i][k] = dp[i][k - 1];
      int j = i + K;
      if (j >= (int)g.size())
        j = g.size() - 1;
      dp[i][k] = min(dp[i][k], dp[j][k - 1]);
    }
  }
}

int lca(int u, int v) {
  u = id[u];
  v = id[v];
  if (u > v)
    swap(u, v);
  int dist = v - u + 1;
  int pos = u;
  P res = P((1 << 26), 0);
  for (int I = 25; I >= 0; I--) {
    int i = (1 << I);
    if (dist < i)
      continue;
    dist -= i;
    res = min(res, dp[pos][I]);
    pos += i;
  }
  return g[res.second];
}

int main() {
  A.init();
  B.init();

  cin >> N >> Tc;

  for (int i = 0; i < N - 1; i++) {
    scanf("%d %d", &a, &b);
    G[a].push_back(b);
    G[b].push_back(a);
  }

  lca_init();

  //  if((int)g.size()>=(1<<21))return 0;

  for (int i = 0; i < Tc; i++) {
    scanf("%d %d %d", &a, &b, &c);
    if (a == 1) {
      int L = id[b], R = id2[b];

      if (al[L] < ar[R - 1] + 1)
        A.add(al[L], ar[R - 1] + 1, c, 0, 0, (1 << 19));
      if (bl[L] < br[R - 1] + 1)
        B.add(bl[L], br[R - 1] + 1, -c, 0, 0, (1 << 19));
    } else {
      int d = lca(b, c);

      ll res1 = 0;
      res1 += A.sum(0, ar[id2[b] - 1] + 1, 0, 0, (1 << 19));
      res1 += B.sum(0, br[id2[b] - 1] + 1, 0, 0, (1 << 19));

      ll res2 = 0;
      res2 += A.sum(0, ar[id2[c] - 1] + 1, 0, 0, (1 << 19));
      res2 += B.sum(0, br[id2[c] - 1] + 1, 0, 0, (1 << 19));

      ll res3 = 0;
      res3 += A.sum(0, ar[id2[d] - 1] + 1, 0, 0, (1 << 19));
      res3 += B.sum(0, br[id2[d] - 1] + 1, 0, 0, (1 << 19));

      cout << res1 + res2 - res3 - res3 << endl;
    }
  }
  return 0;
}

Error: MLE

Fixed submission:

#include <algorithm>
#include <cstdio>
#include <iostream>
#include <queue>
#include <stack>
#include <vector>
using namespace std;
#define MAX_N 300005
#define MAX_M (1 << 21)
#define INF (1 << 28)

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

struct seg {
  ll t[MAX_M];
  ll u[MAX_M];
  void init() {
    for (int i = 0; i < MAX_M; i++)
      t[i] = u[i] = 0;
  }
  void add(int a, int b, ll value, int x, int l, int r) {
    if (b <= l || r <= a) {
      return;
    } else if (a <= l && r <= b) {
      t[x] += value;
      return;
    } else {
      add(a, b, value, x * 2 + 1, l, (l + r) / 2);
      add(a, b, value, x * 2 + 2, (l + r) / 2, r);
      u[x] += value * (min(r, b) - max(l, a));
    }
  }
  ll sum(int a, int b, int x, int l, int r) {
    if (b <= l || r <= a) {
      return 0;
    } else if (a <= l && r <= b) {
      return t[x] * (r - l) + u[x];
    } else {
      ll res = (min(r, b) - max(l, a)) * t[x];
      res += sum(a, b, x * 2 + 1, l, (l + r) / 2);
      res += sum(a, b, x * 2 + 2, (l + r) / 2, r);
      return res;
    }
  }
};

seg A, B;

int N, Tc;
int a, b, c;
vector<int> G[MAX_N];
int parent[MAX_N];
int depth[MAX_N];
int id[MAX_N];
int id2[MAX_N];
vector<int> g;
vector<int> al, ar;
vector<int> bl, br;
int ca = 0, cb = 0;
P dp[MAX_M][26];
stack<int> Q;
void lca_init() {
  for (int i = 0; i < MAX_N; i++)
    parent[i] = -1, id[i] = INF;

  depth[0] = 0;
  Q.push(0);

  while (!Q.empty()) {
    int pos = Q.top();
    Q.pop();
    id[pos] = g.size();
    g.push_back(pos);
    int size = G[pos].size();
    if (size && G[pos][size - 1] == parent[pos])
      G[pos].pop_back(), size--;
    if (size == 0)
      continue;
    Q.push(pos);
    Q.push(G[pos][size - 1]);
    depth[G[pos][size - 1]] = depth[pos] + 1;
    parent[G[pos][size - 1]] = pos;
    G[pos].pop_back();
  }

  /*
  for(int i=0;i<(int)g.size();i++){
    if(i){
      if(depth[g[i-1]]<depth[g[i]])cout<<'+';
      else cout<<'-';
    }
    cout<<g[i];
  }
  cout<<endl;
  */
  al.resize(g.size());
  ar.resize(g.size());
  bl.resize(g.size());
  br.resize(g.size());
  for (int i = 0; i < (int)g.size(); i++) {
    al[i] = bl[i] = INF;
    ar[i] = br[i] = -INF;
    id[g[i]] = min(id[g[i]], i);
    id2[g[i]] = i;
  }

  for (int i = 0; i + 1 < (int)g.size(); i++) {
    if (depth[g[i]] < depth[g[i + 1]]) {
      // cout<<' '<<ca;
      al[i] = ar[i] = ca;
      ca++;
    } else {
      // cout<<' '<<cb;
      bl[i] = br[i] = cb;
      cb++;
    }
  }
  // cout<<endl;

  for (int i = 1; i + 1 < (int)g.size(); i++) {
    ar[i] = max(ar[i], ar[i - 1]);
    br[i] = max(br[i], br[i - 1]);
  }
  for (int i = (int)g.size() - 2; i >= 0; i--) {
    al[i] = min(al[i], al[i + 1]);
    bl[i] = min(bl[i], bl[i + 1]);
  }

  // for(int i=0;i<(int)g.size();i++)cout<<' '<<al[i];cout<<endl;
  //  for(int i=0;i<(int)g.size();i++)cout<<' '<<ar[i];cout<<endl;
  // for(int i=0;i<(int)g.size();i++)cout<<' '<<bl[i];cout<<endl;
  // for(int i=0;i<(int)g.size();i++)cout<<' '<<br[i];;cout<<endl;

  for (int i = 0; i < (int)g.size(); i++)
    dp[i][0] = P(depth[g[i]], i);

  for (int k = 1, K = 1; k < 26; k++, K *= 2) {
    for (int i = 0; i < (int)g.size(); i++) {
      dp[i][k] = dp[i][k - 1];
      int j = i + K;
      if (j >= (int)g.size())
        j = g.size() - 1;
      dp[i][k] = min(dp[i][k], dp[j][k - 1]);
    }
  }
}

int lca(int u, int v) {
  u = id[u];
  v = id[v];
  if (u > v)
    swap(u, v);
  int dist = v - u + 1;
  int pos = u;
  P res = P((1 << 26), 0);
  for (int I = 25; I >= 0; I--) {
    int i = (1 << I);
    if (dist < i)
      continue;
    dist -= i;
    res = min(res, dp[pos][I]);
    pos += i;
  }
  return g[res.second];
}

int main() {
  A.init();
  B.init();

  cin >> N >> Tc;

  for (int i = 0; i < N - 1; i++) {
    scanf("%d %d", &a, &b);
    G[a].push_back(b);
    G[b].push_back(a);
  }

  lca_init();

  //  if((int)g.size()>=(1<<21))return 0;

  for (int i = 0; i < Tc; i++) {
    scanf("%d %d %d", &a, &b, &c);
    if (a == 1) {
      int L = id[b], R = id2[b];

      if (al[L] < ar[R - 1] + 1)
        A.add(al[L], ar[R - 1] + 1, c, 0, 0, (1 << 20));
      if (bl[L] < br[R - 1] + 1)
        B.add(bl[L], br[R - 1] + 1, -c, 0, 0, (1 << 20));
    } else {
      int d = lca(b, c);

      ll res1 = 0;
      res1 += A.sum(0, ar[id2[b] - 1] + 1, 0, 0, (1 << 20));
      res1 += B.sum(0, br[id2[b] - 1] + 1, 0, 0, (1 << 20));

      ll res2 = 0;
      res2 += A.sum(0, ar[id2[c] - 1] + 1, 0, 0, (1 << 20));
      res2 += B.sum(0, br[id2[c] - 1] + 1, 0, 0, (1 << 20));

      ll res3 = 0;
      res3 += A.sum(0, ar[id2[d] - 1] + 1, 0, 0, (1 << 20));
      res3 += B.sum(0, br[id2[d] - 1] + 1, 0, 0, (1 << 20));

      cout << res1 + res2 - res3 - res3 << endl;
    }
  }
  return 0;
}