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

#include <algorithm>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <stack>
#include <vector>
using namespace std;

#define INF (1145141919)
struct Seg {
  long long lazy;
  long long sum;
  long long ml, mm, mr;
  long long maximum;
  Seg() {
    lazy = INF;
    sum = 0;
    ml = mm = mr = 0;
    maximum = 0;
  }
};
inline void output(vector<Seg> &seg, int k, int a, int b) {
  cout << "[k:" << k << ", a:" << a << ", b:" << b << "]"
       << " ";
  cout << "sum:" << seg[k].sum << ", ml:" << seg[k].ml << ", mm:" << seg[k].mm
       << ", mr:" << seg[k].mr << ", maximum:" << seg[k].maximum << endl;
}

inline void lazy(vector<Seg> &seg, int lazyval, int k, int a, int b) {
  seg[k].lazy = lazyval;
  seg[k].sum = seg[k].lazy * (b - a + 1);
  seg[k].mm = seg[k].ml = seg[k].mr = max(seg[k].sum, 0ll);
  seg[k].maximum = lazyval;
}

inline void spread(vector<Seg> &seg, int k, int a, int b) {
  if (seg[k].lazy == INF)
    return;
  seg[k].sum = seg[k].lazy * (b - a + 1);
  if (a != b) {
    int m = (a + b) / 2;
    lazy(seg, seg[k].lazy, k * 2, a, m);
    lazy(seg, seg[k].lazy, k * 2 + 1, m + 1, b);
  }
  seg[k].lazy = INF;
}

inline void reflesh(vector<Seg> &seg, int k, int a, int b) {
  if (a != b) {
    seg[k].sum = seg[k * 2].sum + seg[k * 2 + 1].sum;
    seg[k].ml = max(seg[k * 2].ml, seg[k * 2].sum + seg[k * 2 + 1].ml);
    seg[k].mr = max(seg[k * 2 + 1].mr, seg[k * 2].mr + seg[k * 2 + 1].sum);
    seg[k].mm = max(max(seg[k * 2].mm, seg[k * 2 + 1].mm),
                    seg[k * 2].mr + seg[k * 2 + 1].ml);
    seg[k].maximum = max(seg[k * 2].maximum, seg[k * 2 + 1].maximum);
  }
}

int change(vector<Seg> &seg, int l, int r, int v, int k, int a, int b) {
  spread(seg, k, a, b);
  if (r < a || b < l) {
    return 0;
  }

  if (l <= a && b <= r) {
    lazy(seg, v, k, a, b);
    return 0;
  }
  int m = (a + b) / 2;
  change(seg, l, r, v, k * 2, a, m);
  change(seg, l, r, v, k * 2 + 1, m + 1, b);
  reflesh(seg, k, a, b);
}

struct Three {
  long long ml, mm, mr, sum, maximum;
  Three() : ml(0), mm(0), mr(0), sum(0) {}
  Three(long long ml, long long mm, long long mr, long long sum,
        long long maximum)
      : ml(ml), mm(mm), mr(mr), sum(sum), maximum(maximum) {}
};

inline void output(Three x) {
  // cout << "[k:" << k << ", a:" << a << ", b:" << b << "]" << " ";
  cout << "sum:" << x.sum << ", ml:" << x.ml << ", mm:" << x.mm
       << ", mr:" << x.mr << ", maximum:" << x.maximum << endl;
}

Three getMaximum(vector<Seg> &seg, int l, int r, int k, int a, int b) {
  spread(seg, k, a, b);
  if (r < a || b < l) {
    return Three(0, 0, 0, 0, -INF);
  }
  if (l <= a && b <= r) {
    // output(seg,k,a,b);
    return Three(seg[k].ml, seg[k].mm, seg[k].mr, seg[k].sum, seg[k].maximum);
  }
  int m = (a + b) / 2;
  Three le = getMaximum(seg, l, r, k * 2, a, m);
  Three ri = getMaximum(seg, l, r, k * 2 + 1, m + 1, b);
  Three ans = Three(
      max(ri.ml + le.sum, le.ml), max(max(le.mm, ri.mm), le.mr + ri.ml),
      max(le.mr + ri.sum, ri.mr), ri.sum + le.sum, max(le.maximum, ri.maximum));

  // cout << "(" << a << "," << b << ")" << "[" << ans.ml << " " << ans.mm << "
  // " << ans.mr << " " << ans.sum << "]" << endl;
  reflesh(seg, k, a, b);
  return ans;
}

typedef unsigned int uint;
int in() {
  int x = 0, c;
  for (; (uint)((c = getchar()) - '0') >= 10;) {
    if (c == '-')
      return -in();
    if (!~c)
      throw ~0;
  }
  do {
    x = (x << 3) + (x << 1) + (c - '0');
  } while ((uint)((c = getchar()) - '0') < 10);
  return x;
}

int parent[19][200010];
int depth[200010];
vector<int> gr[200010];
int dfs(int x_, int p_) {
  stack<pair<int, int>> S;
  S.push(make_pair(x_, p_));
  while (S.size()) {
    int x = S.top().first;
    int p = S.top().second;
    S.pop();
    parent[0][x] = p;
    if (p != -1)
      depth[x] = depth[p] + 1;
    for (int i = 0; i < gr[x].size(); i++)
      if (gr[x][i] != p)
        S.push(make_pair(gr[x][i], x));
  }
}

void genlca() {
  for (int i = 1; i < 19; i++) {
    for (int j = 0; j < 200010; j++) {
      if (parent[i - 1][j] != -1)
        parent[i][j] = parent[i - 1][parent[i - 1][j]];
    }
  }
}
pair<int, pair<int, int>> lca(int x, int y) {
  // if( depth[x] > depth[y] ) swap(x,y);
  int d = depth[y] - depth[x];
  for (int i = 0; i < 19; i++) {
    if (d >> i & 1)
      y = parent[i][y];
  }
  if (x == y)
    return make_pair(x, make_pair(-1, -1));

  // if( depth[x] != depth[y] ) cout << "^^" << endl;
  for (int i = 17; i >= 0; i--) {
    if (parent[i][x] != parent[i][y]) {
      x = parent[i][x];
      y = parent[i][y];
    }
  }
  return make_pair(parent[0][x], make_pair(x, y));
}

int gNum = 0;
int belong[200010];
int height[200010];
int lightEdge[200010];
int heavyHeight[200010];
int size[200010];
vector<Seg> seg[200010];

int initval[200010];

int build(int x_, int where) {
  stack<pair<int, int>> S;
  S.push(make_pair(x_, -1));
  vector<int> tmp;
  while (S.size()) {
    int x = S.top().first;
    int prev = S.top().second;
    S.pop();
    if (x == -1)
      continue;
    if (prev != -1 && size[prev] * 2 < size[x]) {
      lightEdge[where] = x;
      continue;
    }
    if (belong[x] != -1) {
      lightEdge[where] = x;
      continue;
    }
    tmp.push_back(x);
    belong[x] = where;
    S.push(make_pair(parent[0][x], x));
  }
  for (int i = 0; i < tmp.size(); i++)
    height[tmp[i]] = tmp.size() - i;
  return tmp.size() + 1;
}

void query1(int st, int en, int val) {
  if (belong[st] != belong[en]) {
    int s = belong[st];
    int a = height[st];
    change(seg[s], 1, a, val, 1, 1, seg[s].size() / 2);
    query1(lightEdge[s], en, val);
  } else {
    int s = belong[st];
    int a = height[st];
    int b = height[en];
    change(seg[s], b, a, val, 1, 1, seg[s].size() / 2);
  }
}

Three query2(int st, int en, Three ri) {
  if (belong[st] != belong[en]) {
    int s = belong[st];
    int a = height[st];
    Three le = getMaximum(seg[s], 1, a, 1, 1, seg[s].size() / 2);
    Three ans =
        Three(max(ri.ml + le.sum, le.ml), max(max(le.mm, ri.mm), le.mr + ri.ml),
              max(le.mr + ri.sum, ri.mr), ri.sum + le.sum,
              max(le.maximum, ri.maximum));
    return query2(lightEdge[s], en, ans);

  } else {
    int s = belong[st];
    int a = height[st];
    int b = height[en];
    Three le = getMaximum(seg[s], b, a, 1, 1, seg[s].size() / 2);
    Three ans =
        Three(max(ri.ml + le.sum, le.ml), max(max(le.mm, ri.mm), le.mr + ri.ml),
              max(le.mr + ri.sum, ri.mr), ri.sum + le.sum,
              max(le.maximum, ri.maximum));
    return ans;
  }
}

int main() {

  memset(parent, -1, sizeof(parent));
  memset(belong, -1, sizeof(belong));
  memset(lightEdge, -1, sizeof(lightEdge));
  int n = in();
  int m = in();
  for (int i = 0; i < n; i++)
    initval[i] = in();
  for (int i = 0; i < n - 1; i++) {
    int a = in();
    int b = in();
    a--, b--;
    gr[a].push_back(b);
    gr[b].push_back(a);
  }
  dfs(0, -1);
  genlca();
  vector<pair<int, int>> sorted;
  for (int i = 0; i < n; i++) {
    sorted.push_back(make_pair(-depth[i], i));
  }
  sort(sorted.begin(), sorted.end());
  for (int j = 0; j < n; j++) {
    int i = sorted[j].second;
    size[i]++;
    if (parent[0][i] != -1)
      size[parent[0][i]] += size[i];
  }
  for (int j = 0; j < n; j++) {
    int i = sorted[j].second;
    if (belong[i] == -1) {
      int sz = build(i, gNum) - 1;
      int sz_ = 1;
      while (sz_ < sz)
        sz_ *= 2;
      seg[gNum].resize(2 * sz_);
      gNum++;
    }
  }
  for (int i = 0; i < n; i++) {
    change(seg[belong[i]], height[i], height[i], initval[i], 1, 1,
           seg[belong[i]].size() / 2);
  }
  for (int i = 0; i < m; i++) {
    int t, a, b, c;
    t = in();
    a = in();
    b = in();
    c = in();
    a--, b--;
    if (depth[a] > depth[b])
      swap(a, b);
    pair<int, pair<int, int>> tmp = lca(a, b);
    int p = tmp.first;
    int pa = tmp.first;
    int pb = tmp.second.second;
    if (t == 1) {
      if (p == a) {
        query1(b, p, c);
      } else {
        query1(a, p, c);
        query1(b, pb, c);
      }
    } else {
      Three res, ri, le;
      if (p == a) {
        res = query2(b, p, Three(0, 0, 0, 0, -INF));
      } else {
        le = query2(a, p, Three(0, 0, 0, 0, -INF));
        ri = query2(b, pb, Three(0, 0, 0, 0, -INF));
        res = Three(-1, max(max(le.mm, ri.mm), le.ml + ri.ml), -1, -1,
                    max(ri.maximum, le.maximum));
      }
      if (res.maximum < 0)
        printf("%lld\\n", res.maximum);
      else
        printf("%lld\\n", res.mm);
    }
  }
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <stack>
#include <vector>
using namespace std;

#define INF (1145141919)
struct Seg {
  long long lazy;
  long long sum;
  long long ml, mm, mr;
  long long maximum;
  Seg() {
    lazy = INF;
    sum = 0;
    ml = mm = mr = 0;
    maximum = 0;
  }
};
inline void output(vector<Seg> &seg, int k, int a, int b) {
  cout << "[k:" << k << ", a:" << a << ", b:" << b << "]"
       << " ";
  cout << "sum:" << seg[k].sum << ", ml:" << seg[k].ml << ", mm:" << seg[k].mm
       << ", mr:" << seg[k].mr << ", maximum:" << seg[k].maximum << endl;
}

inline void lazy(vector<Seg> &seg, int lazyval, int k, int a, int b) {
  seg[k].lazy = lazyval;
  seg[k].sum = seg[k].lazy * (b - a + 1);
  seg[k].mm = seg[k].ml = seg[k].mr = max(seg[k].sum, 0ll);
  seg[k].maximum = lazyval;
}

inline void spread(vector<Seg> &seg, int k, int a, int b) {
  if (seg[k].lazy == INF)
    return;
  seg[k].sum = seg[k].lazy * (b - a + 1);
  if (a != b) {
    int m = (a + b) / 2;
    lazy(seg, seg[k].lazy, k * 2, a, m);
    lazy(seg, seg[k].lazy, k * 2 + 1, m + 1, b);
  }
  seg[k].lazy = INF;
}

inline void reflesh(vector<Seg> &seg, int k, int a, int b) {
  if (a != b) {
    seg[k].sum = seg[k * 2].sum + seg[k * 2 + 1].sum;
    seg[k].ml = max(seg[k * 2].ml, seg[k * 2].sum + seg[k * 2 + 1].ml);
    seg[k].mr = max(seg[k * 2 + 1].mr, seg[k * 2].mr + seg[k * 2 + 1].sum);
    seg[k].mm = max(max(seg[k * 2].mm, seg[k * 2 + 1].mm),
                    seg[k * 2].mr + seg[k * 2 + 1].ml);
    seg[k].maximum = max(seg[k * 2].maximum, seg[k * 2 + 1].maximum);
  }
}

int change(vector<Seg> &seg, int l, int r, int v, int k, int a, int b) {
  spread(seg, k, a, b);
  if (r < a || b < l) {
    return 0;
  }

  if (l <= a && b <= r) {
    lazy(seg, v, k, a, b);
    return 0;
  }
  int m = (a + b) / 2;
  change(seg, l, r, v, k * 2, a, m);
  change(seg, l, r, v, k * 2 + 1, m + 1, b);
  reflesh(seg, k, a, b);
}

struct Three {
  long long ml, mm, mr, sum, maximum;
  Three() : ml(0), mm(0), mr(0), sum(0) {}
  Three(long long ml, long long mm, long long mr, long long sum,
        long long maximum)
      : ml(ml), mm(mm), mr(mr), sum(sum), maximum(maximum) {}
};

inline void output(Three x) {
  // cout << "[k:" << k << ", a:" << a << ", b:" << b << "]" << " ";
  cout << "sum:" << x.sum << ", ml:" << x.ml << ", mm:" << x.mm
       << ", mr:" << x.mr << ", maximum:" << x.maximum << endl;
}

Three getMaximum(vector<Seg> &seg, int l, int r, int k, int a, int b) {
  spread(seg, k, a, b);
  if (r < a || b < l) {
    return Three(0, 0, 0, 0, -INF);
  }
  if (l <= a && b <= r) {
    // output(seg,k,a,b);
    return Three(seg[k].ml, seg[k].mm, seg[k].mr, seg[k].sum, seg[k].maximum);
  }
  int m = (a + b) / 2;
  Three le = getMaximum(seg, l, r, k * 2, a, m);
  Three ri = getMaximum(seg, l, r, k * 2 + 1, m + 1, b);
  Three ans = Three(
      max(ri.ml + le.sum, le.ml), max(max(le.mm, ri.mm), le.mr + ri.ml),
      max(le.mr + ri.sum, ri.mr), ri.sum + le.sum, max(le.maximum, ri.maximum));

  // cout << "(" << a << "," << b << ")" << "[" << ans.ml << " " << ans.mm << "
  // " << ans.mr << " " << ans.sum << "]" << endl;
  reflesh(seg, k, a, b);
  return ans;
}

typedef unsigned int uint;
int in() {
  int x = 0, c;
  for (; (uint)((c = getchar()) - '0') >= 10;) {
    if (c == '-')
      return -in();
    if (!~c)
      throw ~0;
  }
  do {
    x = (x << 3) + (x << 1) + (c - '0');
  } while ((uint)((c = getchar()) - '0') < 10);
  return x;
}

int parent[19][200010];
int depth[200010];
vector<int> gr[200010];
int dfs(int x_, int p_) {
  stack<pair<int, int>> S;
  S.push(make_pair(x_, p_));
  while (S.size()) {
    int x = S.top().first;
    int p = S.top().second;
    S.pop();
    parent[0][x] = p;
    if (p != -1)
      depth[x] = depth[p] + 1;
    for (int i = 0; i < gr[x].size(); i++)
      if (gr[x][i] != p)
        S.push(make_pair(gr[x][i], x));
  }
}

void genlca() {
  for (int i = 1; i < 19; i++) {
    for (int j = 0; j < 200010; j++) {
      if (parent[i - 1][j] != -1)
        parent[i][j] = parent[i - 1][parent[i - 1][j]];
    }
  }
}
pair<int, pair<int, int>> lca(int x, int y) {
  // if( depth[x] > depth[y] ) swap(x,y);
  int d = depth[y] - depth[x];
  for (int i = 0; i < 19; i++) {
    if (d >> i & 1)
      y = parent[i][y];
  }
  if (x == y)
    return make_pair(x, make_pair(-1, -1));

  // if( depth[x] != depth[y] ) cout << "^^" << endl;
  for (int i = 17; i >= 0; i--) {
    if (parent[i][x] != parent[i][y]) {
      x = parent[i][x];
      y = parent[i][y];
    }
  }
  return make_pair(parent[0][x], make_pair(x, y));
}

int gNum = 0;
int belong[200010];
int height[200010];
int heavyEdge[200010];
int heavyHeight[200010];

vector<Seg> seg[200010];

int initval[200010];

int build(int x_, int where) {
  stack<int> S;
  S.push(x_);
  vector<int> tmp;
  while (S.size()) {
    int x = S.top();
    S.pop();
    if (x == -1)
      continue;
    if (belong[x] != -1) {
      heavyEdge[where] = x;
      continue;
    }
    tmp.push_back(x);
    belong[x] = where;
    S.push(parent[0][x]);
  }
  for (int i = 0; i < tmp.size(); i++)
    height[tmp[i]] = tmp.size() - i;
  return tmp.size() + 1;
}

void query1(int st, int en, int val) {
  if (belong[st] != belong[en]) {
    int s = belong[st];
    int a = height[st];
    change(seg[s], 1, a, val, 1, 1, seg[s].size() / 2);
    query1(heavyEdge[s], en, val);
  } else {

    int s = belong[st];
    int a = height[st];
    int b = height[en];
    // cout << s << endl;
    change(seg[s], b, a, val, 1, 1, seg[s].size() / 2);
  }
}

Three query2(int st, int en, Three ri) {
  if (belong[st] != belong[en]) {
    int s = belong[st];
    int a = height[st];
    Three le = getMaximum(seg[s], 1, a, 1, 1, seg[s].size() / 2);
    Three ans =
        Three(max(ri.ml + le.sum, le.ml), max(max(le.mm, ri.mm), le.mr + ri.ml),
              max(le.mr + ri.sum, ri.mr), ri.sum + le.sum,
              max(le.maximum, ri.maximum));
    /*output(le);
    output(ri);
    output(ans);
    puts("=---=");
    */
    return query2(heavyEdge[s], en, ans);

  } else {
    int s = belong[st];
    int a = height[st];
    int b = height[en];
    Three le = getMaximum(seg[s], b, a, 1, 1, seg[s].size() / 2);
    Three ans =
        Three(max(ri.ml + le.sum, le.ml), max(max(le.mm, ri.mm), le.mr + ri.ml),
              max(le.mr + ri.sum, ri.mr), ri.sum + le.sum,
              max(le.maximum, ri.maximum));
    /*output(le);
    output(ri);
    output(ans);
    puts("=---=");
    */

    // cout << "<" << st << "," << en << ">" << a << " " <<  b << " < < <" <<
    // endl; output(ans);
    return ans;
  }
}

int main() {
  memset(parent, -1, sizeof(parent));
  memset(belong, -1, sizeof(belong));
  memset(heavyEdge, -1, sizeof(heavyEdge));
  int n = in();
  int m = in();
  for (int i = 0; i < n; i++)
    initval[i] = in();
  for (int i = 0; i < n - 1; i++) {
    int a = in();
    int b = in();
    a--, b--;
    gr[a].push_back(b);
    gr[b].push_back(a);
  }
  dfs(0, -1);
  genlca();
  vector<pair<int, int>> sorted;
  for (int i = 0; i < n; i++) {
    sorted.push_back(make_pair(-depth[i], i));
  }
  sort(sorted.begin(), sorted.end());
  for (int j = 0; j < n; j++) {
    int i = sorted[j].second;
    if (belong[i] == -1) {
      int sz = build(i, gNum) - 1;
      int sz_ = 1;
      while (sz_ < sz)
        sz_ *= 2;
      seg[gNum].resize(2 * sz_);
      gNum++;
    }
  }
  for (int i = 0; i < n; i++) {
    change(seg[belong[i]], height[i], height[i], initval[i], 1, 1,
           seg[belong[i]].size() / 2);
  }
  /*
  for(int i = 0 ; i < gNum ; i++){
          for(int j = 1 ; j <= seg[i].size()/2 ; j++){
                  cout << getMaximum(seg[i],j,j,1,1,seg[i].size()/2).maximum <<
  " ";
          }
          cout << "[" << heavyEdge[i] << "]";
          cout << endl;
  }
  */
  for (int i = 0; i < m; i++) {
    int t, a, b, c;
    t = in();
    a = in();
    b = in();
    c = in();
    a--, b--;
    if (depth[a] > depth[b])
      swap(a, b);
    pair<int, pair<int, int>> tmp = lca(a, b);
    int p = tmp.first;
    int pa = tmp.first;
    int pb = tmp.second.second;
    if (t == 1) {
      if (p == a) {
        // cout << belong[b] << " " << belong[a] << endl;
        query1(b, p, c);
      } else {
        query1(a, p, c);
        query1(b, pb, c);
      }
    } else {
      // cout << a+1 << "-" << p+1 << " " << p+1 << "-" <<  b+1 << "(p:" << p+1
      // << ")" << endl;
      Three res, ri, le;
      if (p == a) {
        res = query2(b, p, Three(0, 0, 0, 0, -INF));
      } else {
        // cout << a+1 << " " << pa+1 << " " << pb+1 << " " << b+1 << endl;
        le = query2(a, p, Three(0, 0, 0, 0, -INF));
        ri = query2(b, pb, Three(0, 0, 0, 0, -INF));
        // output(le);
        // output(ri);
        res = Three(-1, max(max(le.mm, ri.mm), le.ml + ri.ml), -1, -1,
                    max(ri.maximum, le.maximum));
      }
      if (res.maximum < 0)
        printf("%lld\\n", res.maximum);
      else
        printf("%lld\\n", res.mm);
      // output(res);
    }
  }
  // cout << lca(1,2) << endl;
}