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

#include <bits/stdc++.h>
using namespace std;
typedef long long int64;
const static int64 INF = 1LL << 58;
typedef vector<vector<int>> Graph;
enum Type { MODIFICATION, OUTPUT };
// ??????????????°??§???????????§??????????????????The ?????§??????

struct Query {
  Type type;
  int a, b, c;
};
struct Node {
  int ParIdx, ParDepth, Deep, Size;
};
struct SegNode {
  int64 Ans, Sum, Left, Right, Lazy;
};
inline SegNode Merge(SegNode a, SegNode b) {
  SegNode ret;
  ret.Ans = max(max(a.Ans, b.Ans), a.Right + b.Left);
  ret.Sum = max(-INF, a.Sum + b.Sum);
  ret.Left = max(a.Sum + b.Left, a.Left);
  ret.Right = max(b.Sum + a.Right, b.Right);
  ret.Lazy = INF;
  return (ret);
}
struct SegmentTree {
  vector<SegNode> data;
  int sz;
  SegmentTree(size_t n) {
    sz = 1;
    while (sz < n)
      sz <<= 1;
    data.assign(2 * sz + 1, (SegNode){0, 0, 0, 0, INF});
  }
  inline void Update(int a, int b, int64 value, int k, int l, int r) {
    if (a >= r || b <= l)
      return;
    if (a <= l && r <= b) {
      data[k].Lazy = value;
      data[k].Left = data[k].Right = data[k].Ans =
          (value >= 0 ? (r - l) * value : value);
      data[k].Sum = (r - l) * value;
      return;
    }
    LazyMarge(k, l, r);
    Update(a, b, value, 2 * k + 1, l, (l + r) >> 1);
    Update(a, b, value, 2 * k + 2, (l + r) >> 1, r);
    data[k] = Merge(data[2 * k + 1], data[2 * k + 2]);
  }

  inline void LazyMarge(int k, int l, int r) {
    if (data[k].Lazy == INF)
      return;
    Update(l, (l + r) >> 1, data[k].Lazy, 2 * k + 1, l, (l + r) >> 1);
    Update((l + r) >> 1, r, data[k].Lazy, 2 * k + 2, (l + r) >> 1, r);
    data[k].Lazy = INF;
  }
  void Update(int a, int b, int64 x) { Update(a, b, x, 0, 0, sz); }
  void init() {
    for (int k = sz - 2; k >= 0; k--) {
      data[k] = Merge(data[2 * k + 1], data[2 * k + 2]);
    }
  }
  inline SegNode Query(int a, int b, int k, int l, int r) {
    if (a >= r || b <= l)
      return ((SegNode){-INF, 0, -INF, -INF});
    if (a <= l && r <= b)
      return (data[k]);
    LazyMarge(k, l, r);
    return (Merge(Query(a, b, 2 * k + 1, l, (l + r) >> 1),
                  Query(a, b, 2 * k + 2, (l + r) >> 1, r)));
  }
  SegNode Query(int a, int b) { return (Query(a, b, 0, 0, sz)); }
};

int N, Q, Value[200000];
Graph graph;
vector<Query> querys;
int sz[200000];
Node pool[200000];
int ptr = 0, nowTreeIdx[200000], nowTreeDepth[200000];
vector<int> nodes[200000];
vector<SegmentTree> segs;

int sizeDFS() {
  stack<pair<int, int>> st;
  st.push(make_pair(N / 2, -1));
  while (!st.empty()) {
    pair<int, int> p = st.top();
    st.pop();
    if (p.first < 0) {
      p.first = -p.first - 1;
      for (int i = 0; i < graph[p.first].size(); i++) {
        if (graph[p.first][i] != p.second) {
          sz[p.first] += sz[graph[p.first][i]];
        }
      }
    } else {
      st.push(make_pair(-p.first - 1, p.second));
      for (int i = 0; i < graph[p.first].size(); i++) {
        if (graph[p.first][i] != p.second) {
          st.push(make_pair(graph[p.first][i], p.first));
        }
      }
    }
  }
}

queue<pair<pair<int, int>, pair<int, int>>> que;
inline void makeDFS() {
  pool[ptr++] = (Node){-1, -1, 0, 0};
  que.push(make_pair(make_pair(N / 2, -1), make_pair(0, 0)));
  while (!que.empty()) {
    pair<pair<int, int>, pair<int, int>> p = que.front();
    que.pop();
    int treeIdx = p.second.first, treeDepth = p.second.second;
    int idx = p.first.first, prev = p.first.second;
    pool[treeIdx].Size++;
    nodes[treeIdx].push_back(idx);
    nowTreeIdx[idx] = treeIdx;
    nowTreeDepth[idx] = treeDepth;
    int ret = 0;
    for (int i = 0; i < graph[idx].size(); i++) {
      if (graph[idx][i] != prev)
        ret = max(ret, sz[graph[idx][i]]);
    }
    for (int i = 0; i < graph[idx].size(); i++) {
      if (graph[idx][i] != prev) {
        if (ret == sz[graph[idx][i]]) {
          que.push(make_pair(make_pair(graph[idx][i], idx),
                             make_pair(treeIdx, treeDepth + 1)));
          ret = -1;
        } else {
          pool[ptr++] = (Node){treeIdx, treeDepth, pool[treeIdx].Deep + 1, 0};
          que.push(
              make_pair(make_pair(graph[idx][i], idx), make_pair(ptr - 1, 0)));
        }
      }
    }
  }
}
void CentroidPathDecomposition() {
  sizeDFS();
  makeDFS();
}

void Running() {
  CentroidPathDecomposition();
  for (int i = 0; i < ptr; i++) {
    segs.push_back(SegmentTree(pool[i].Size));
    for (int j = 0; j < nodes[i].size(); j++) {
      SegNode *node = &(segs[i].data[j + segs[i].sz - 1]);
      node->Left = node->Right = node->Ans = node->Sum = Value[nodes[i][j]];
    }
    segs[i].init();
  }

  for (int i = 0; i < Q; i++) {
    const Query &q = querys[i];
    if (q.type == MODIFICATION) {
      int tidxA = nowTreeIdx[q.a], tidxB = nowTreeIdx[q.b];
      int tdepA = nowTreeDepth[q.a], tdepB = nowTreeDepth[q.b];
      while (tidxA != tidxB) {
        if (pool[tidxA].Deep > pool[tidxB].Deep) {
          segs[tidxA].Update(0, tdepA + 1, q.c);
          tdepA = pool[tidxA].ParDepth;
          tidxA = pool[tidxA].ParIdx;
        } else {
          segs[tidxB].Update(0, tdepB + 1, q.c);
          tdepB = pool[tidxB].ParDepth;
          tidxB = pool[tidxB].ParIdx;
        }
      }
      if (tdepA > tdepB)
        swap(tdepA, tdepB);
      segs[tidxA].Update(tdepA, tdepB + 1, q.c);
    } else {

      SegNode aa = (SegNode){-INF, 0, -INF, -INF};
      SegNode bb = (SegNode){-INF, 0, -INF, -INF};
      int tidxA = nowTreeIdx[q.a], tidxB = nowTreeIdx[q.b];
      int tdepA = nowTreeDepth[q.a], tdepB = nowTreeDepth[q.b];
      while (tidxA != tidxB) {
        if (pool[tidxA].Deep > pool[tidxB].Deep) {
          aa = Merge(segs[tidxA].Query(0, tdepA + 1), aa);
          tdepA = pool[tidxA].ParDepth;
          tidxA = pool[tidxA].ParIdx;
        } else {
          bb = Merge(segs[tidxB].Query(0, tdepB + 1), bb);
          tdepB = pool[tidxB].ParDepth;
          tidxB = pool[tidxB].ParIdx;
        }
      }
      if (tdepA < tdepB) {
        swap(aa.Left, aa.Right);
        aa = Merge(aa, segs[tidxA].Query(tdepA, tdepB + 1));
        aa = Merge(aa, bb);
        printf("%lld\\n", aa.Ans);
      } else {
        swap(bb.Left, bb.Right);
        bb = Merge(bb, segs[tidxA].Query(tdepB, tdepA + 1));
        bb = Merge(bb, aa);
        printf("%lld\\n", bb.Ans);
      }
    }
  }
}

void Input() {
  scanf("%d %d", &N, &Q);
  for (int i = 0; i < N; i++) {
    scanf("%d", &Value[i]);
  }
  graph.resize(N);
  for (int i = 0; i < N - 1; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    --a, --b;
    graph[a].push_back(b);
    graph[b].push_back(a);
  }
  querys.resize(Q);
  for (int i = 0; i < Q; i++) {
    int t, a, b, c;
    scanf("%d %d %d %d", &t, &a, &b, &c);
    --a, --b;
    if (t == 1) {
      querys[i] = (Query){MODIFICATION, a, b, c};
    } else {
      querys[i] = (Query){OUTPUT, a, b, c};
    }
  }
}

int main() {
  Input();
  Running();
  return (0);
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long int64;
const static int64 INF = 1LL << 58;
typedef vector<vector<int>> Graph;
enum Type { MODIFICATION, OUTPUT };
// ??????????????°??§???????????§??????????????????The ?????§??????

struct Query {
  Type type;
  int a, b, c;
};
struct Node {
  int ParIdx, ParDepth, Deep, Size;
};
struct SegNode {
  int64 Ans, Sum, Left, Right, Lazy;
};
SegNode Merge(SegNode a, SegNode b) {
  SegNode ret;
  ret.Ans = max(max(a.Ans, b.Ans), a.Right + b.Left);
  ret.Sum = max(-INF, a.Sum + b.Sum);
  ret.Left = max(a.Sum + b.Left, a.Left);
  ret.Right = max(b.Sum + a.Right, b.Right);
  ret.Lazy = INF;
  return (ret);
}
struct SegmentTree {
  vector<SegNode> data;
  int sz;
  SegmentTree(size_t n) {
    sz = 1;
    while (sz < n)
      sz <<= 1;
    data.assign(2 * sz + 1, (SegNode){0, 0, 0, 0, INF});
  }
  void Update(int a, int b, int64 value, int k, int l, int r) {
    if (a >= r || b <= l)
      return;
    if (a <= l && r <= b) {
      data[k].Lazy = value;
      data[k].Left = data[k].Right = data[k].Ans =
          (value >= 0 ? (r - l) * value : value);
      data[k].Sum = (r - l) * value;
      return;
    }
    LazyMarge(k, l, r);
    Update(a, b, value, 2 * k + 1, l, (l + r) >> 1);
    Update(a, b, value, 2 * k + 2, (l + r) >> 1, r);
    data[k] = Merge(data[2 * k + 1], data[2 * k + 2]);
  }

  void LazyMarge(int k, int l, int r) {
    if (data[k].Lazy == INF)
      return;
    Update(l, (l + r) >> 1, data[k].Lazy, 2 * k + 1, l, (l + r) >> 1);
    Update((l + r) >> 1, r, data[k].Lazy, 2 * k + 2, (l + r) >> 1, r);
    data[k].Lazy = INF;
  }
  void Update(int a, int b, int64 x) { Update(a, b, x, 0, 0, sz); }
  void init() {
    for (int k = sz - 2; k >= 0; k--) {
      data[k] = Merge(data[2 * k + 1], data[2 * k + 2]);
    }
  }
  SegNode Query(int a, int b, int k, int l, int r) {
    if (a >= r || b <= l)
      return ((SegNode){-INF, 0, -INF, -INF});
    if (a <= l && r <= b)
      return (data[k]);
    LazyMarge(k, l, r);
    return (Merge(Query(a, b, 2 * k + 1, l, (l + r) >> 1),
                  Query(a, b, 2 * k + 2, (l + r) >> 1, r)));
  }
  SegNode Query(int a, int b) { return (Query(a, b, 0, 0, sz)); }
};

int N, Q, Value[200000];
Graph graph;
vector<Query> querys;
int sz[200000];
Node pool[200000];
int ptr = 0, nowTreeIdx[200000], nowTreeDepth[200000];
vector<int> nodes[200000];
vector<SegmentTree> segs;

int sizeDFS() {
  stack<pair<int, int>> st;
  st.push(make_pair(N / 2, -1));
  while (!st.empty()) {
    pair<int, int> p = st.top();
    st.pop();
    if (p.first < 0) {
      p.first = -p.first - 1;
      for (int i = 0; i < graph[p.first].size(); i++) {
        if (graph[p.first][i] != p.second) {
          sz[p.first] += sz[graph[p.first][i]];
        }
      }
    } else {
      st.push(make_pair(-p.first - 1, p.second));
      for (int i = 0; i < graph[p.first].size(); i++) {
        if (graph[p.first][i] != p.second) {
          st.push(make_pair(graph[p.first][i], p.first));
        }
      }
    }
  }
}

queue<pair<pair<int, int>, pair<int, int>>> que;
inline void makeDFS() {
  pool[ptr++] = (Node){-1, -1, 0, 0};
  que.push(make_pair(make_pair(N / 2, -1), make_pair(0, 0)));
  while (!que.empty()) {
    pair<pair<int, int>, pair<int, int>> p = que.front();
    que.pop();
    int treeIdx = p.second.first, treeDepth = p.second.second;
    int idx = p.first.first, prev = p.first.second;
    pool[treeIdx].Size++;
    nodes[treeIdx].push_back(idx);
    nowTreeIdx[idx] = treeIdx;
    nowTreeDepth[idx] = treeDepth;
    int ret = 0;
    for (int i = 0; i < graph[idx].size(); i++) {
      if (graph[idx][i] != prev)
        ret = max(ret, sz[graph[idx][i]]);
    }
    for (int i = 0; i < graph[idx].size(); i++) {
      if (graph[idx][i] != prev) {
        if (ret == sz[graph[idx][i]]) {
          que.push(make_pair(make_pair(graph[idx][i], idx),
                             make_pair(treeIdx, treeDepth + 1)));
          ret = -1;
        } else {
          pool[ptr++] = (Node){treeIdx, treeDepth, pool[treeIdx].Deep + 1, 0};
          que.push(
              make_pair(make_pair(graph[idx][i], idx), make_pair(ptr - 1, 0)));
        }
      }
    }
  }
}
void CentroidPathDecomposition() {
  sizeDFS();
  makeDFS();
}

void Running() {
  CentroidPathDecomposition();
  for (int i = 0; i < ptr; i++) {
    segs.push_back(SegmentTree(pool[i].Size));
    for (int j = 0; j < nodes[i].size(); j++) {
      SegNode *node = &(segs[i].data[j + segs[i].sz - 1]);
      node->Left = node->Right = node->Ans = node->Sum = Value[nodes[i][j]];
    }
    segs[i].init();
  }

  for (int i = 0; i < Q; i++) {
    const Query &q = querys[i];
    if (q.type == MODIFICATION) {
      int tidxA = nowTreeIdx[q.a], tidxB = nowTreeIdx[q.b];
      int tdepA = nowTreeDepth[q.a], tdepB = nowTreeDepth[q.b];
      while (tidxA != tidxB) {
        if (pool[tidxA].Deep > pool[tidxB].Deep) {
          segs[tidxA].Update(0, tdepA + 1, q.c);
          tdepA = pool[tidxA].ParDepth;
          tidxA = pool[tidxA].ParIdx;
        } else {
          segs[tidxB].Update(0, tdepB + 1, q.c);
          tdepB = pool[tidxB].ParDepth;
          tidxB = pool[tidxB].ParIdx;
        }
      }
      if (tdepA > tdepB)
        swap(tdepA, tdepB);
      segs[tidxA].Update(tdepA, tdepB + 1, q.c);
    } else {

      SegNode aa = (SegNode){-INF, 0, -INF, -INF};
      SegNode bb = (SegNode){-INF, 0, -INF, -INF};
      int tidxA = nowTreeIdx[q.a], tidxB = nowTreeIdx[q.b];
      int tdepA = nowTreeDepth[q.a], tdepB = nowTreeDepth[q.b];
      while (tidxA != tidxB) {
        if (pool[tidxA].Deep > pool[tidxB].Deep) {
          aa = Merge(segs[tidxA].Query(0, tdepA + 1), aa);
          tdepA = pool[tidxA].ParDepth;
          tidxA = pool[tidxA].ParIdx;
        } else {
          bb = Merge(segs[tidxB].Query(0, tdepB + 1), bb);
          tdepB = pool[tidxB].ParDepth;
          tidxB = pool[tidxB].ParIdx;
        }
      }
      if (tdepA < tdepB) {
        swap(aa.Left, aa.Right);
        aa = Merge(aa, segs[tidxA].Query(tdepA, tdepB + 1));
        aa = Merge(aa, bb);
        printf("%lld\\n", aa.Ans);
      } else {
        swap(bb.Left, bb.Right);
        bb = Merge(bb, segs[tidxA].Query(tdepB, tdepA + 1));
        bb = Merge(bb, aa);
        printf("%lld\\n", bb.Ans);
      }
    }
  }
}

void Input() {
  scanf("%d %d", &N, &Q);
  for (int i = 0; i < N; i++) {
    scanf("%d", &Value[i]);
  }
  graph.resize(N);
  for (int i = 0; i < N - 1; i++) {
    int a, b;
    scanf("%d %d", &a, &b);
    --a, --b;
    graph[a].push_back(b);
    graph[b].push_back(a);
  }
  querys.resize(Q);
  for (int i = 0; i < Q; i++) {
    int t, a, b, c;
    scanf("%d %d %d %d", &t, &a, &b, &c);
    --a, --b;
    if (t == 1) {
      querys[i] = (Query){MODIFICATION, a, b, c};
    } else {
      querys[i] = (Query){OUTPUT, a, b, c};
    }
  }
}

int main() {
  Input();
  Running();
  return (0);
}