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

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

struct HLDecomposition {
  int n, pos;
  vector<vector<int>> G;
  vector<int> vid, head, sub, hvy, par, dep, inv, type;

  HLDecomposition() {}
  HLDecomposition(int sz)
      : n(sz), pos(0), G(n), vid(n, -1), head(n), sub(n, 1), hvy(n, -1), par(n),
        dep(n), inv(n), type(n) {}

  void add_edge(int u, int v) {
    G[u].push_back(v);
    G[v].push_back(u);
  }

  void build(vector<int> rs = {0}) {
    int c = 0;
    for (int r : rs) {
      dfs(r);
      bfs(r, c++);
    }
  }

  void dfs(int rt) {
    using T = pair<int, int>;
    stack<T> st;
    par[rt] = -1;
    dep[rt] = 0;
    st.emplace(rt, 0);
    while (!st.empty()) {
      int v = st.top().first;
      int &i = st.top().second;
      if (i < (int)G[v].size()) {
        int u = G[v][i++];
        if (u == par[v])
          continue;
        par[u] = v;
        dep[u] = dep[v] + 1;
        st.emplace(u, 0);
      } else {
        st.pop();
        int res = 0;
        for (int u : G[v]) {
          if (u == par[v])
            continue;
          sub[v] += sub[u];
          if (res < sub[u])
            res = sub[u], hvy[v] = u;
        }
      }
    }
  }

  void bfs(int r, int c) {
    int &k = pos;
    queue<int> q({r});
    while (!q.empty()) {
      int h = q.front();
      q.pop();
      for (int i = h; i != -1; i = hvy[i]) {
        type[i] = c;
        vid[i] = k++;
        inv[vid[i]] = i;
        head[i] = h;
        for (int j : G[i])
          if (j != par[i] && j != hvy[i])
            q.push(j);
      }
    }
  }

  // for_each(vertex)
  // [l,r] <- attention!!
  void for_each(int u, int v, const function<void(int, int)> &f) {
    while (1) {
      if (vid[u] > vid[v])
        swap(u, v);
      f(max(vid[head[v]], vid[u]), vid[v]);
      if (head[u] != head[v])
        v = par[head[v]];
      else
        break;
    }
  }

  // for_each(edge)
  // [l,r] <- attention!!
  void for_each_edge(int u, int v, const function<void(int, int)> &f) {
    while (1) {
      if (vid[u] > vid[v])
        swap(u, v);
      if (head[u] != head[v]) {
        f(vid[head[v]], vid[v]);
        v = par[head[v]];
      } else {
        if (u != v)
          f(vid[u] + 1, vid[v]);
        break;
      }
    }
  }

  int lca(int u, int v) {
    while (1) {
      if (vid[u] > vid[v])
        swap(u, v);
      if (head[u] == head[v])
        return u;
      v = par[head[v]];
    }
  }

  int distance(int u, int v) { return dep[u] + dep[v] - 2 * dep[lca(u, v)]; }
};

template <typename T, typename E> struct SegmentTree {
  using F = function<T(T, T)>;
  using G = function<T(T, E)>;
  using H = function<E(E, E)>;
  using P = function<E(E, size_t)>;
  int n;
  F f;
  G g;
  H h;
  T ti;
  E ei;
  P p;
  vector<T> dat;
  vector<E> laz;
  SegmentTree(
      int n_, F f, G g, H h, T ti, E ei,
      P p =
          [](E a, size_t b) {
            b++;
            return a;
          })
      : f(f), g(g), h(h), ti(ti), ei(ei), p(p) {
    init(n_);
  }
  void init(int n_) {
    n = 1;
    while (n < n_)
      n *= 2;
    dat.assign(2 * n - 1, ti);
    laz.assign(2 * n - 1, ei);
  }
  void build(int n_, vector<T> v) {
    for (int i = 0; i < n_; i++)
      dat[i + n - 1] = v[i];
    for (int i = n - 2; i >= 0; i--)
      dat[i] = f(dat[i * 2 + 1], dat[i * 2 + 2]);
  }
  inline void eval(int len, int k) {
    if (laz[k] == ei)
      return;
    if (k * 2 + 1 < n * 2 - 1) {
      laz[k * 2 + 1] = h(laz[k * 2 + 1], laz[k]);
      laz[k * 2 + 2] = h(laz[k * 2 + 2], laz[k]);
    }
    dat[k] = g(dat[k], p(laz[k], len));
    laz[k] = ei;
  }
  T update(int a, int b, E x, int k, int l, int r) {
    eval(r - l, k);
    if (r <= a || b <= l)
      return dat[k];
    if (a <= l && r <= b) {
      laz[k] = h(laz[k], x);
      return g(dat[k], p(laz[k], r - l));
    }
    return dat[k] = f(update(a, b, x, k * 2 + 1, l, (l + r) / 2),
                      update(a, b, x, k * 2 + 2, (l + r) / 2, r));
  }
  T update(int a, int b, E x) { return update(a, b, x, 0, 0, n); }
  T query(int a, int b, int k, int l, int r) {
    eval(r - l, k);
    if (r <= a || b <= l)
      return ti;
    if (a <= l && r <= b)
      return dat[k];
    T vl = query(a, b, k * 2 + 1, l, (l + r) / 2);
    T vr = query(a, b, k * 2 + 2, (l + r) / 2, r);
    return f(vl, vr);
  }
  T query(int a, int b) { return query(a, b, 0, 0, n); }
  void update(int k, T x) {
    query(k, k + 1); // evaluate
    k += n - 1;
    dat[k] = x;
    while (k) {
      k = (k - 1) / 2;
      dat[k] = f(dat[k * 2 + 1], dat[k * 2 + 2]);
    }
  }
};

// INSERT ABOVE HERE
signed main() {
  int n, q;
  scanf("%d %d", &n, &q);
  HLDecomposition hld(n);
  vector<int> w(n);
  for (int i = 0; i < n; i++)
    cin >> w[i];
  using P = pair<int, int>;
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    hld.add_edge(a, b);
  }
  hld.build();

  using T = tuple<int, int, int, int, int, int, int>;

  T d1(-1, -1, -1, -1, -1, -1, -1);
  P d0(-1, -114514);

  auto &par = hld.par;
  auto &vid = hld.vid;
  ;
  auto con = [&](int a, int b) { return par[a] == b || par[b] == a; };

  auto f = [&](T a, T b) {
    if (a > b)
      swap(a, b);

    if (get<0>(a) < 0)
      return b;
    if (con(get<0>(a), get<1>(b)))
      swap(a, b);

    int al, ar, as, ava, avi, avl, avr;
    tie(al, ar, as, ava, avi, avl, avr) = a;
    int bl, br, bs, bva, bvi, bvl, bvr;
    tie(bl, br, bs, bva, bvi, bvl, bvr) = b;

    if (!con(ar, bl)) {
      if (con(ar, br)) {
        swap(bl, br);
        swap(bvl, bvr);
      } else if (con(al, bl)) {
        swap(al, ar);
        swap(avl, avr);
      } else {
        return d1;
      }
    }
    int cl = al, cr = br, cs = as + bs;
    int cva = ava + bva, cvi = max(avi, bvi), cvl = avl, cvr = bvr;
    cvi = max(cvi, avr + bvl);
    cvl = max(cvl, ava + bvl);
    cvr = max(cvr, avr + bva);

    return T(cl, cr, cs, cva, cvi, cvl, cvr);
  };

  auto g = [&](T a, P p) {
    if (p == d0)
      return a;
    int al, ar, as, ava, avi, avl, avr;
    tie(al, ar, as, ava, avi, avl, avr) = a;
    int v = p.first, b = p.second;
    if (~v)
      al = ar = v, as = 1;
    if (b >= 0)
      return T(al, ar, as, b * as, b * as, b * as, b * as);
    return T(al, ar, as, b * as, b, b, b);
  };

  auto h = [&](P a, P b) {
    a.first++;
    return b;
  };

  vector<T> vt(n);
  for (int i = 0; i < n; i++)
    vt[vid[i]] = g(d1, P(i, w[i]));

  SegmentTree<T, P> seg(n, f, g, h, d1, d0);
  seg.build(n, vt);

  while (q--) {
    int t, a, b, c;
    scanf("%d %d %d %d", &t, &a, &b, &c);
    a--;
    b--;
    if (t == 1) {
      hld.for_each(a, b, [&](int l, int r) { seg.update(l, r + 1, P(-1, c)); });
    }
    if (t == 2) {
      vector<T> v;
      hld.for_each(a, b, [&](int l, int r) {
        T x = seg.query(l, r + 1);
        for (int i = 0; i < (int)v.size(); i++) {
          T y = f(v[i], x);
          if (~get<0>(y)) {
            v[i] = y;
            return;
          }
        }
        v.push_back(x);
      });
      while (v.size() > 1) {
        T x = v.back();
        v.pop_back();
        for (int i = 0; i < (int)v.size(); i++) {
          T y = f(v[i], x);
          if (~get<0>(y)) {
            v[i] = y;
            break;
          }
        }
      }
      int vl, vr, vs, vva, vvi, vvl, vvr;
      tie(vl, vr, vs, vva, vvi, vvl, vvr) = v.front();
      printf("%d\\n", max({vva, vvi, vvl, vvr}));
    }
  }

  return 0;
}


Error: TLE

Fixed submission:

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

struct HLDecomposition {
  int n, pos;
  vector<vector<int>> G;
  vector<int> vid, head, sub, hvy, par, dep, inv, type, ps, pt;

  HLDecomposition() {}
  HLDecomposition(int sz)
      : n(sz), pos(0), G(n), vid(n, -1), head(n), sub(n, 1), hvy(n, -1), par(n),
        dep(n), inv(n), type(n), ps(n), pt(n) {}

  void add_edge(int u, int v) {
    G[u].push_back(v);
    G[v].push_back(u);
  }

  void build(vector<int> rs = {0}) {
    int c = 0;
    for (int r : rs) {
      dfs(r);
      bfs(r, c++);
    }
  }

  void dfs(int rt) {
    using T = pair<int, int>;
    stack<T> st;
    par[rt] = -1;
    dep[rt] = 0;
    st.emplace(rt, 0);
    while (!st.empty()) {
      int v = st.top().first;
      int &i = st.top().second;
      if (i < (int)G[v].size()) {
        int u = G[v][i++];
        if (u == par[v])
          continue;
        par[u] = v;
        dep[u] = dep[v] + 1;
        st.emplace(u, 0);
      } else {
        st.pop();
        for (int j = 0; j < (int)G[v].size(); j++) {
          int &u = G[v][j];
          if (u == par[v]) {
            swap(u, G[v].back());
            continue;
          }
          sub[v] += sub[u];
          if (sub[u] > sub[G[v].front()])
            swap(u, G[v].front());
        }
      }
    }
  }

  void bfs(int r, int c) {
    using T = tuple<int, int, int>;
    stack<T> st;
    st.emplace(r, r, 0);
    while (!st.empty()) {
      int v, h;
      tie(v, h, ignore) = st.top();
      int &i = get<2>(st.top());
      if (!i) {
        type[v] = c;
        ps[v] = vid[v] = pos++;
        inv[vid[v]] = v;
        head[v] = h;
        hvy[v] = (G[v].empty() ? -1 : G[v][0]);
        if (hvy[v] == par[v])
          hvy[v] = -1;
      }
      if (i < (int)G[v].size()) {
        int u = G[v][i++];
        if (u == par[v])
          continue;
        st.emplace(u, (hvy[v] == u ? h : u), 0);
      } else {
        st.pop();
        pt[v] = pos;
      }
    }
  }

  // for_each(vertex)
  // [l,r] <- attention!!
  void for_each(int u, int v, const function<void(int, int)> &f) {
    while (1) {
      if (vid[u] > vid[v])
        swap(u, v);
      f(max(vid[head[v]], vid[u]), vid[v]);
      if (head[u] != head[v])
        v = par[head[v]];
      else
        break;
    }
  }

  // for_each(edge)
  // [l,r] <- attention!!
  void for_each_edge(int u, int v, const function<void(int, int)> &f) {
    while (1) {
      if (vid[u] > vid[v])
        swap(u, v);
      if (head[u] != head[v]) {
        f(vid[head[v]], vid[v]);
        v = par[head[v]];
      } else {
        if (u != v)
          f(vid[u] + 1, vid[v]);
        break;
      }
    }
  }

  int lca(int u, int v) {
    while (1) {
      if (vid[u] > vid[v])
        swap(u, v);
      if (head[u] == head[v])
        return u;
      v = par[head[v]];
    }
  }

  int distance(int u, int v) { return dep[u] + dep[v] - 2 * dep[lca(u, v)]; }
};

template <typename T, typename E> struct SegmentTree {
  using F = function<T(T, T)>;
  using G = function<T(T, E)>;
  using H = function<E(E, E)>;
  using P = function<E(E, size_t)>;
  int n;
  F f;
  G g;
  H h;
  T ti;
  E ei;
  P p;
  vector<T> dat;
  vector<E> laz;
  SegmentTree(
      int n_, F f, G g, H h, T ti, E ei,
      P p =
          [](E a, size_t b) {
            b++;
            return a;
          })
      : f(f), g(g), h(h), ti(ti), ei(ei), p(p) {
    init(n_);
  }
  void init(int n_) {
    n = 1;
    while (n < n_)
      n *= 2;
    dat.assign(2 * n - 1, ti);
    laz.assign(2 * n - 1, ei);
  }
  void build(int n_, vector<T> v) {
    for (int i = 0; i < n_; i++)
      dat[i + n - 1] = v[i];
    for (int i = n - 2; i >= 0; i--)
      dat[i] = f(dat[i * 2 + 1], dat[i * 2 + 2]);
  }
  inline void eval(int len, int k) {
    if (laz[k] == ei)
      return;
    if (k * 2 + 1 < n * 2 - 1) {
      laz[k * 2 + 1] = h(laz[k * 2 + 1], laz[k]);
      laz[k * 2 + 2] = h(laz[k * 2 + 2], laz[k]);
    }
    dat[k] = g(dat[k], p(laz[k], len));
    laz[k] = ei;
  }
  T update(int a, int b, E x, int k, int l, int r) {
    eval(r - l, k);
    if (r <= a || b <= l)
      return dat[k];
    if (a <= l && r <= b) {
      laz[k] = h(laz[k], x);
      return g(dat[k], p(laz[k], r - l));
    }
    return dat[k] = f(update(a, b, x, k * 2 + 1, l, (l + r) / 2),
                      update(a, b, x, k * 2 + 2, (l + r) / 2, r));
  }
  T update(int a, int b, E x) { return update(a, b, x, 0, 0, n); }
  T query(int a, int b, int k, int l, int r) {
    eval(r - l, k);
    if (r <= a || b <= l)
      return ti;
    if (a <= l && r <= b)
      return dat[k];
    T vl = query(a, b, k * 2 + 1, l, (l + r) / 2);
    T vr = query(a, b, k * 2 + 2, (l + r) / 2, r);
    return f(vl, vr);
  }
  T query(int a, int b) { return query(a, b, 0, 0, n); }
  void update(int k, T x) {
    query(k, k + 1); // evaluate
    k += n - 1;
    dat[k] = x;
    while (k) {
      k = (k - 1) / 2;
      dat[k] = f(dat[k * 2 + 1], dat[k * 2 + 2]);
    }
  }
};

// INSERT ABOVE HERE
signed main() {
  int n, q;
  scanf("%d %d", &n, &q);
  HLDecomposition hld(n);
  vector<int> w(n);
  for (int i = 0; i < n; i++)
    cin >> w[i];
  using P = pair<int, int>;
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    cin >> a >> b;
    a--;
    b--;
    hld.add_edge(a, b);
  }
  hld.build();

  using T = tuple<int, int, int, int, int, int, int>;

  T d1(-1, -1, -1, -1, -1, -1, -1);
  P d0(-1, -114514);

  auto &par = hld.par;
  auto &vid = hld.vid;
  ;
  auto con = [&](int a, int b) { return par[a] == b || par[b] == a; };

  auto f = [&](T a, T b) {
    if (a > b)
      swap(a, b);

    if (get<0>(a) < 0)
      return b;
    if (con(get<0>(a), get<1>(b)))
      swap(a, b);

    int al, ar, as, ava, avi, avl, avr;
    tie(al, ar, as, ava, avi, avl, avr) = a;
    int bl, br, bs, bva, bvi, bvl, bvr;
    tie(bl, br, bs, bva, bvi, bvl, bvr) = b;

    if (!con(ar, bl)) {
      if (con(ar, br)) {
        swap(bl, br);
        swap(bvl, bvr);
      } else if (con(al, bl)) {
        swap(al, ar);
        swap(avl, avr);
      } else {
        return d1;
      }
    }
    int cl = al, cr = br, cs = as + bs;
    int cva = ava + bva, cvi = max(avi, bvi), cvl = avl, cvr = bvr;
    cvi = max(cvi, avr + bvl);
    cvl = max(cvl, ava + bvl);
    cvr = max(cvr, avr + bva);

    return T(cl, cr, cs, cva, cvi, cvl, cvr);
  };

  auto g = [&](T a, P p) {
    if (p == d0)
      return a;
    int al, ar, as, ava, avi, avl, avr;
    tie(al, ar, as, ava, avi, avl, avr) = a;
    int v = p.first, b = p.second;
    if (~v)
      al = ar = v, as = 1;
    if (b >= 0)
      return T(al, ar, as, b * as, b * as, b * as, b * as);
    return T(al, ar, as, b * as, b, b, b);
  };

  auto h = [&](P a, P b) {
    a.first++;
    return b;
  };

  vector<T> vt(n);
  for (int i = 0; i < n; i++)
    vt[vid[i]] = g(d1, P(i, w[i]));

  SegmentTree<T, P> seg(n, f, g, h, d1, d0);
  seg.build(n, vt);

  while (q--) {
    int t, a, b, c;
    scanf("%d %d %d %d", &t, &a, &b, &c);
    a--;
    b--;
    if (t == 1) {
      hld.for_each(a, b, [&](int l, int r) { seg.update(l, r + 1, P(-1, c)); });
    }
    if (t == 2) {
      vector<T> v;
      hld.for_each(a, b, [&](int l, int r) {
        T x = seg.query(l, r + 1);
        for (int i = 0; i < (int)v.size(); i++) {
          T y = f(v[i], x);
          if (~get<0>(y)) {
            v[i] = y;
            return;
          }
        }
        v.push_back(x);
      });
      while (v.size() > 1) {
        T x = v.back();
        v.pop_back();
        for (int i = 0; i < (int)v.size(); i++) {
          T y = f(v[i], x);
          if (~get<0>(y)) {
            v[i] = y;
            break;
          }
        }
      }
      int vl, vr, vs, vva, vvi, vvl, vvr;
      tie(vl, vr, vs, vva, vvi, vvl, vvr) = v.front();
      printf("%d\\n", max({vva, vvi, vvl, vvr}));
    }
  }

  return 0;
}