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

#include <algorithm>
#include <cassert>
#include <cctype>
#include <climits>
#include <cstdlib>
#include <cstring>
#include <map>
#include <queue>
#include <stdio.h>
#include <unistd.h>
#include <vector>

using namespace std;
typedef long long ll;
#define int ll
int n, q;
int w[222222];
vector<int> g[222222];
vector<int> ch[222222];
vector<int> G[222222];
int tdepth[222222];
int f[222222];
int par[222222];
int ix[222222];
int chi[222222];
int sz[222222];
int depth[222222];
int dpar[222222];
bool used[222222];
int kk;
const pair<int, int> F(-114514810, -114514810);

void dfs(int v, int pre) {
  queue<pair<int, int>> q;
  queue<int> lf;
  q.push(make_pair(v, pre));
  par[v] = -1;
  tdepth[v] = 0;
  for (int i = 0; i < n; i++)
    sz[i] = 1;
  while (!q.empty()) {
    pair<int, int> p = std::move(q.front());
    q.pop();
    bool o = false;
    for (int i = 0; i < int(g[p.first].size()); i++) {
      if (p.second != g[p.first][i]) {
        par[g[p.first][i]] = p.first;
        tdepth[g[p.first][i]] = tdepth[p.first] + 1;
        ch[p.first].emplace_back(g[p.first][i]);
        q.emplace(g[p.first][i], p.first);
        o = true;
      }
    }
    if (!o) {
      lf.push(p.first);
      used[p.first] = true;
    }
  }
  while (!lf.empty()) {
    int p = lf.front();
    lf.pop();
    sz[par[p]] += sz[p];
    if (par[p] >= 0 && !used[par[p]]) {
      used[par[p]] = true;
      lf.emplace(par[p]);
    }
  }
}

void dec(int v, int d) {
  queue<tuple<int, int, int>> q;
  q.push(make_tuple(v, d, kk));
  while (!q.empty()) {
    auto pp = std::move(q.front());
    v = get<0>(pp);
    d = get<1>(pp);
    int K = get<2>(pp);
    q.pop();
    ix[v] = G[K].size();
    G[K].emplace_back(v);
    f[v] = K;
    depth[K] = d;
    int idx = -1;
    for (int i = 0; i < int(ch[v].size()); i++) {
      if (idx == -1 || sz[ch[v][i]] > /*= sz[v]/2+1*/ sz[idx]) {
        idx = ch[v][i];
      }
    }
    if (idx >= 0) {
      q.emplace(idx, d, K);
    }
    for (int i = 0; i < int(ch[v].size()); i++) {
      if (ch[v][i] != idx) {
        kk += 1;
        chi[kk] = ch[v][i];
        dpar[kk] = f[v];
        q.emplace(ch[v][i], d + 1, kk);
      }
    }
  }
}

struct segtree {
  int n;
  struct dat {
    int ans, all, left, right;
    dat() {}
    dat(const dat &l, const dat &r) {
      ans = max(max(l.ans, r.ans), l.right + r.left);
      all = l.all + r.all;
      left = max(l.left, l.all + r.left);
      right = max(r.right, r.all + l.right);
    }
    explicit dat(int a) { ans = all = left = right = a; }
  };
  vector<dat> d;
  vector<pair<int, int>> lazy;
  segtree() {}
  void init(vector<int> &g) {
    n = 1;
    while (n < int(g.size()))
      n <<= 1;
    d.resize((n << 1) - 1);
    lazy.resize((n << 1) - 1, F);
    for (int i = 0; i < int(g.size()); i++) {
      d[i + (n - 1)] = std::move(dat(w[g[i]]));
    }
    for (int i = n - 2; i >= 0; i--) {
      d[i] = std::move(dat(d[(i << 1) + 1], d[(i << 1) + 2]));
    }
  }
  void push(int k) {
    if (lazy[k] != F) {
      d[k].all = lazy[k].first;
      if (lazy[k].first >= 0) {
        d[k].left = lazy[k].first;
        d[k].right = lazy[k].first;
        d[k].ans = d[k].all;
      } else {
        d[k].left = lazy[k].second;
        d[k].right = lazy[k].second;
        d[k].ans = lazy[k].second;
      }
      if (k * 2 + 2 < int(d.size())) {
        lazy[k].first >>= 1;
        lazy[k * 2 + 1] = lazy[k * 2 + 2] = std::move(lazy[k]);
      }
      lazy[k] = F;
    }
  }
  inline void update(int i, int j, int l, int r, int k, int v) {
    push(k);
    if (j <= l || r <= i)
      return;
    if (l <= i && j <= r) {
      lazy[k] = std::move(make_pair(v, v / (j - i)));
      push(k);
    } else {
      update(i, (i + j) >> 1, l, r, (k << 1) + 1, v >> 1);
      update((i + j) >> 1, j, l, r, (k << 1) + 2, v >> 1);
      d[k] = std::move(dat(d[(k << 1) + 1], d[(k << 1) + 2]));
    }
  }
  inline void update(int l, int r, int v) { update(0, n, l, r, 0, v * n); }
  inline dat get(int i, int j, int l, int r, int k) {
    // printf("%d %d %d %d %d\\n",i,j,l,r,k);
    static const int inf = 1000000011;
    if (j <= l || r <= i)
      return std::move(dat(-inf));
    push(k);
    if (l <= i && j <= r) {
      return d[k];
    } else {
      dat ld = std::move(get(i, (i + j) >> 1, l, r, (k << 1) + 1));
      dat rd = std::move(get((i + j) >> 1, j, l, r, (k << 1) + 2));
      if (ld.all == -inf)
        return std::move(rd);
      if (rd.all == -inf)
        return std::move(ld);
      return std::move(dat(ld, rd));
    }
  }
  inline dat get(int l, int r) { return std::move(get(0, n, l, r, 0)); }
  void print() {
    int c = 1, w = n, cnt = 0;
    for (int i = 0; i < 2 * n - 1; i++) {
      push(i);
      cnt += 1;
      printf("(ans:%d,l:%d,r:%d,a:%d) ", d[i].ans, d[i].left, d[i].right,
             d[i].all);
      if (c == cnt) {
        puts("");
        c *= 2;
        cnt = 0;
        w /= 2;
      }
    }
    puts("");
  }
};

segtree t[222222];

const int inf = 1145141919;
segtree::dat DF(-inf);
inline segtree::dat merge(segtree::dat a, segtree::dat b) {
  if (a.all == -inf)
    return b;
  if (b.all == -inf)
    return a;
  return segtree::dat(a, b);
}

inline void update(int a, int b, int c) {
  int aa = f[a];
  int bb = f[b];
  while (depth[aa] > depth[bb]) {
    t[aa].update(0, ix[a] + 1, c);
    a = par[G[aa][0]];
    aa = dpar[aa];
  }
  while (depth[bb] > depth[aa]) {
    t[bb].update(0, ix[b] + 1, c);
    b = par[G[bb][0]];
    bb = dpar[bb];
  }
  while (aa != bb) {
    t[aa].update(0, ix[a] + 1, c);
    t[bb].update(0, ix[b] + 1, c);
    a = par[G[aa][0]];
    b = par[G[bb][0]];
    aa = dpar[aa];
    bb = dpar[bb];
  }
  t[aa].update(min(ix[a], ix[b]), max(ix[a], ix[b]) + 1, c);
}

segtree::dat mid;
segtree::dat dat[514], rev_dat[514];
segtree::dat trev_dat[514];
int dat_size, rev_dat_size;

inline int get(int a, int b) {
  int aa = f[a];
  int bb = f[b];
  segtree::dat al = DF, bl = DF;
  dat_size = rev_dat_size = 0;
  int lc, rc;
  lc = rc = 0;
  while (depth[aa] > depth[bb]) {
    al = merge(std::move(t[aa].get(0, ix[a] + 1)), al);
    a = par[G[aa][0]];
    aa = dpar[aa];
    lc += 1;
  }
  while (depth[bb] > depth[aa]) {
    bl = merge(std::move(t[bb].get(0, ix[b] + 1)), bl);
    b = par[G[bb][0]];
    bb = dpar[bb];
    rc += 1;
  }
  while (aa != bb) {
    al = merge(t[aa].get(0, ix[a] + 1), al);
    bl = merge(std::move(t[bb].get(0, ix[b] + 1)), bl);
    a = par[G[aa][0]];
    b = par[G[bb][0]];
    aa = dpar[aa];
    bb = dpar[bb];
    lc += 1;
  }
  mid = std::move(t[aa].get(min(ix[a], ix[b]), max(ix[a], ix[b]) + 1));
  // fprintf(stderr,"%d %d\\n",aa,lc+rc);
  if (ix[a] > ix[b]) {
    swap(mid.left, mid.right);
  }
  swap(al.left, al.right);
  al = merge(al, mid);
  al = merge(al, bl);
  return al.ans;
}

const int S = 11451419;
char buf[S];
int ccc = 0;

inline int getc() { return buf[ccc++]; }

inline int get_i() {
  int c = getc();
  if (c == '-') {
    c = getc();
    int res = c - '0';
    for (;;) {
      c = getc();
      if (isdigit(c)) {
        res = /*res*10*/ (res << 3) + (res << 1) + c - '0';
      } else {
        break;
      }
    }
    return -res;
  } else {
    int res = c - '0';
    for (;;) {
      c = getc();
      if (isdigit(c)) {
        res = (res << 3) + (res << 1) + c - '0';
      } else {
        break;
      }
    }
    return res;
  }
}

signed main(void) {
  read(0, buf, S);
  n = get_i();
  q = get_i();
  // printf("%d %d\\n",n,q);
  for (int i = 0; i < n; i++) {
    w[i] = get_i();
    // par[i] = -1;
  }
  for (int i = 0; i < n - 1; i++) {
    int s, t;
    s = get_i();
    t = get_i();
    --s;
    --t;
    g[s].emplace_back(t);
    g[t].emplace_back(s);
  }
  memset(par, -1, sizeof(par));
  dfs(0, -1);
  kk = 0;
  dec(0, 0);
  kk += 1;
  for (int i = 0; i < kk; i++) {
    t[i].init(G[i]);
  }
  for (int i = 0; i < q; i++) {
    int a, b, c, d;
    a = get_i();
    b = get_i();
    c = get_i();
    d = get_i();
    --b;
    --c;
    if (a == 1) {
      update(b, c, d);
    } else {
      printf("%lld\\n", get(b, c));
    }
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <cassert>
#include <cctype>
#include <climits>
#include <cstdlib>
#include <cstring>
#include <map>
#include <queue>
#include <stdio.h>
#include <unistd.h>
#include <vector>

using namespace std;
typedef long long ll;
#define int ll
int n, q;
int w[222222];
vector<int> g[222222];
vector<int> ch[222222];
vector<int> G[222222];
int tdepth[222222];
int f[222222];
int par[222222];
int ix[222222];
int chi[222222];
int sz[222222];
int depth[222222];
int dpar[222222];
bool used[222222];
int kk;
const pair<int, int> F(-114514810, -114514810);

void dfs(int v, int pre) {
  queue<pair<int, int>> q;
  queue<int> lf;
  q.push(make_pair(v, pre));
  par[v] = -1;
  tdepth[v] = 0;
  for (int i = 0; i < n; i++)
    sz[i] = 1;
  while (!q.empty()) {
    pair<int, int> p = std::move(q.front());
    q.pop();
    bool o = false;
    for (int i = 0; i < int(g[p.first].size()); i++) {
      if (p.second != g[p.first][i]) {
        par[g[p.first][i]] = p.first;
        tdepth[g[p.first][i]] = tdepth[p.first] + 1;
        ch[p.first].emplace_back(g[p.first][i]);
        q.emplace(g[p.first][i], p.first);
        o = true;
      }
    }
    if (!o) {
      lf.push(p.first);
      used[p.first] = true;
    }
  }
  while (!lf.empty()) {
    int p = lf.front();
    lf.pop();
    if (par[p] >= 0 && !used[par[p]]) {
      sz[par[p]] += sz[p];
      used[par[p]] = true;
      lf.emplace(par[p]);
    }
  }
}

void dec(int v, int d) {
  queue<tuple<int, int, int>> q;
  q.push(make_tuple(v, d, kk));
  while (!q.empty()) {
    auto pp = std::move(q.front());
    v = get<0>(pp);
    d = get<1>(pp);
    int K = get<2>(pp);
    q.pop();
    ix[v] = G[K].size();
    G[K].emplace_back(v);
    f[v] = K;
    depth[K] = d;
    int idx = -1;
    for (int i = 0; i < int(ch[v].size()); i++) {
      if (idx == -1 || sz[ch[v][i]] > /*= sz[v]/2+1*/ sz[idx]) {
        idx = ch[v][i];
      }
    }
    if (idx >= 0) {
      q.emplace(idx, d, K);
    }
    for (int i = 0; i < int(ch[v].size()); i++) {
      if (ch[v][i] != idx) {
        kk += 1;
        chi[kk] = ch[v][i];
        dpar[kk] = f[v];
        q.emplace(ch[v][i], d + 1, kk);
      }
    }
  }
}

struct segtree {
  int n;
  struct dat {
    int ans, all, left, right;
    dat() {}
    dat(const dat &l, const dat &r) {
      ans = max(max(l.ans, r.ans), l.right + r.left);
      all = l.all + r.all;
      left = max(l.left, l.all + r.left);
      right = max(r.right, r.all + l.right);
    }
    explicit dat(int a) { ans = all = left = right = a; }
  };
  vector<dat> d;
  vector<pair<int, int>> lazy;
  segtree() {}
  void init(vector<int> &g) {
    n = 1;
    while (n < int(g.size()))
      n <<= 1;
    d.resize((n << 1) - 1);
    lazy.resize((n << 1) - 1, F);
    for (int i = 0; i < int(g.size()); i++) {
      d[i + (n - 1)] = std::move(dat(w[g[i]]));
    }
    for (int i = n - 2; i >= 0; i--) {
      d[i] = std::move(dat(d[(i << 1) + 1], d[(i << 1) + 2]));
    }
  }
  void push(int k) {
    if (lazy[k] != F) {
      d[k].all = lazy[k].first;
      if (lazy[k].first >= 0) {
        d[k].left = lazy[k].first;
        d[k].right = lazy[k].first;
        d[k].ans = d[k].all;
      } else {
        d[k].left = lazy[k].second;
        d[k].right = lazy[k].second;
        d[k].ans = lazy[k].second;
      }
      if (k * 2 + 2 < int(d.size())) {
        lazy[k].first >>= 1;
        lazy[k * 2 + 1] = lazy[k * 2 + 2] = std::move(lazy[k]);
      }
      lazy[k] = F;
    }
  }
  inline void update(int i, int j, int l, int r, int k, int v) {
    push(k);
    if (j <= l || r <= i)
      return;
    if (l <= i && j <= r) {
      lazy[k] = std::move(make_pair(v, v / (j - i)));
      push(k);
    } else {
      update(i, (i + j) >> 1, l, r, (k << 1) + 1, v >> 1);
      update((i + j) >> 1, j, l, r, (k << 1) + 2, v >> 1);
      d[k] = std::move(dat(d[(k << 1) + 1], d[(k << 1) + 2]));
    }
  }
  inline void update(int l, int r, int v) { update(0, n, l, r, 0, v * n); }
  inline dat get(int i, int j, int l, int r, int k) {
    // printf("%d %d %d %d %d\\n",i,j,l,r,k);
    static const int inf = 1000000011;
    if (j <= l || r <= i)
      return std::move(dat(-inf));
    push(k);
    if (l <= i && j <= r) {
      return d[k];
    } else {
      dat ld = std::move(get(i, (i + j) >> 1, l, r, (k << 1) + 1));
      dat rd = std::move(get((i + j) >> 1, j, l, r, (k << 1) + 2));
      if (ld.all == -inf)
        return std::move(rd);
      if (rd.all == -inf)
        return std::move(ld);
      return std::move(dat(ld, rd));
    }
  }
  inline dat get(int l, int r) { return std::move(get(0, n, l, r, 0)); }
  void print() {
    int c = 1, w = n, cnt = 0;
    for (int i = 0; i < 2 * n - 1; i++) {
      push(i);
      cnt += 1;
      printf("(ans:%d,l:%d,r:%d,a:%d) ", d[i].ans, d[i].left, d[i].right,
             d[i].all);
      if (c == cnt) {
        puts("");
        c *= 2;
        cnt = 0;
        w /= 2;
      }
    }
    puts("");
  }
};

segtree t[222222];

const int inf = 1145141919;
segtree::dat DF(-inf);
inline segtree::dat merge(segtree::dat a, segtree::dat b) {
  if (a.all == -inf)
    return b;
  if (b.all == -inf)
    return a;
  return segtree::dat(a, b);
}

inline void update(int a, int b, int c) {
  int aa = f[a];
  int bb = f[b];
  while (depth[aa] > depth[bb]) {
    t[aa].update(0, ix[a] + 1, c);
    a = par[G[aa][0]];
    aa = dpar[aa];
  }
  while (depth[bb] > depth[aa]) {
    t[bb].update(0, ix[b] + 1, c);
    b = par[G[bb][0]];
    bb = dpar[bb];
  }
  while (aa != bb) {
    t[aa].update(0, ix[a] + 1, c);
    t[bb].update(0, ix[b] + 1, c);
    a = par[G[aa][0]];
    b = par[G[bb][0]];
    aa = dpar[aa];
    bb = dpar[bb];
  }
  t[aa].update(min(ix[a], ix[b]), max(ix[a], ix[b]) + 1, c);
}

segtree::dat mid;
segtree::dat dat[514], rev_dat[514];
segtree::dat trev_dat[514];
int dat_size, rev_dat_size;

inline int get(int a, int b) {
  int aa = f[a];
  int bb = f[b];
  segtree::dat al = DF, bl = DF;
  dat_size = rev_dat_size = 0;
  int lc, rc;
  lc = rc = 0;
  while (depth[aa] > depth[bb]) {
    al = merge(std::move(t[aa].get(0, ix[a] + 1)), al);
    a = par[G[aa][0]];
    aa = dpar[aa];
    lc += 1;
  }
  while (depth[bb] > depth[aa]) {
    bl = merge(std::move(t[bb].get(0, ix[b] + 1)), bl);
    b = par[G[bb][0]];
    bb = dpar[bb];
    rc += 1;
  }
  while (aa != bb) {
    al = merge(t[aa].get(0, ix[a] + 1), al);
    bl = merge(std::move(t[bb].get(0, ix[b] + 1)), bl);
    a = par[G[aa][0]];
    b = par[G[bb][0]];
    aa = dpar[aa];
    bb = dpar[bb];
    lc += 1;
  }
  mid = std::move(t[aa].get(min(ix[a], ix[b]), max(ix[a], ix[b]) + 1));
  // fprintf(stderr,"%d %d\\n",aa,lc+rc);
  if (ix[a] > ix[b]) {
    swap(mid.left, mid.right);
  }
  swap(al.left, al.right);
  al = merge(al, mid);
  al = merge(al, bl);
  return al.ans;
}

const int S = 11451419;
char buf[S];
int ccc = 0;

inline int getc() { return buf[ccc++]; }

inline int get_i() {
  int c = getc();
  if (c == '-') {
    c = getc();
    int res = c - '0';
    for (;;) {
      c = getc();
      if (isdigit(c)) {
        res = /*res*10*/ (res << 3) + (res << 1) + c - '0';
      } else {
        break;
      }
    }
    return -res;
  } else {
    int res = c - '0';
    for (;;) {
      c = getc();
      if (isdigit(c)) {
        res = (res << 3) + (res << 1) + c - '0';
      } else {
        break;
      }
    }
    return res;
  }
}

signed main(void) {
  read(0, buf, S);
  n = get_i();
  q = get_i();
  // printf("%d %d\\n",n,q);
  for (int i = 0; i < n; i++) {
    w[i] = get_i();
    // par[i] = -1;
  }
  for (int i = 0; i < n - 1; i++) {
    int s, t;
    s = get_i();
    t = get_i();
    --s;
    --t;
    g[s].emplace_back(t);
    g[t].emplace_back(s);
  }
  memset(par, -1, sizeof(par));
  dfs(0, -1);
  kk = 0;
  dec(0, 0);
  kk += 1;
  for (int i = 0; i < kk; i++) {
    t[i].init(G[i]);
  }
  for (int i = 0; i < q; i++) {
    int a, b, c, d;
    a = get_i();
    b = get_i();
    c = get_i();
    d = get_i();
    --b;
    --c;
    if (a == 1) {
      update(b, c, d);
    } else {
      printf("%lld\\n", get(b, c));
    }
  }
  return 0;
}