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

#include "bits/stdc++.h"
#include <unordered_map>
#include <unordered_set>
#pragma warning(disable : 4996)
using namespace std;

typedef vector<int> Value1;
typedef vector<short> Value2;
const Value1 Zero1{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1};
const Value2 Zero2{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

struct Node {
  Value1
      sum; //??´??°????????????.
           //?????????????????§??????????????????????????¨????????????????????????????????????.
  Value2 lazy; //?????¶???????????????????????????????????????
  Node() : sum(Zero1) { lazy = Zero2; }
};
int aN, Q;
struct lazy_segtree {
  int N;
  vector<Node> dat;
  lazy_segtree(int n) : N(1) {
    while (N < n)
      N *= 2;
    dat.resize(2 * N);
    for (int t = N; t < N + aN; ++t) {
      dat[t].sum = vector<int>{1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
    }
    for (int i = N - 1; i >= 1; --i) {
      update_node(i);
    }
  }
  void fresh() {
    for (int i = N - 1; i >= 1; --i) {
      update_node(i);
    }
  }

  Value2 lazy_connect(const Value2 l, const Value2 r) {
    Value2 v(l);
    for (int i = 0; i < 11; ++i) {
      v[i] = r[v[i]];
    }
    return v;
  }

  void lazy_func(const int k, const int a, const int b) {
    Value1 next(11);
    for (int i = 0; i < 11; ++i) {
      next[dat[k].lazy[i]] += dat[k].sum[i];
    }
    dat[k].sum = next;
  }

  Value1 connect(const Value1 l, const Value1 r) {
    Value1 next(l);
    for (int i = 0; i < 11; ++i)
      next[i] += r[i];
    return next;
  }

  // inline??????????????¨??§???????????¨????????????!(??????)
  inline void lazy_evaluate_node(int k, int a, int b) {
    lazy_func(k, a, b);
    if (k <
        N) { // 2*k(???????????????) < 2*N (???????????°) ?????????????????§.
             // ???????????????????????????????????¬???????????¨??????.
      dat[2 * k].lazy = lazy_connect(
          dat[2 * k].lazy,
          dat[k].lazy); //?¬???????????????¬??????????????£?????????.
      dat[2 * k + 1].lazy = lazy_connect(dat[2 * k + 1].lazy, dat[k].lazy);
    }
    dat[k].lazy = Zero2;
  }

  inline void update_node(
      int k) { // k???????????¢?????????????????????????????¨?????????.
               // ????????\\????????¨????????????????????????????????????????????????????????§if???????????????.
    dat[k].sum = connect(dat[2 * k].sum, dat[2 * k + 1].sum);
  }

  // update(l,r,v) := [l,r)?????´??°??????. ?????????0-indexed.
  void update(int l, int r, const Value2 &v, int k = 1, int a = 0, int b = -1) {
    if (b == -1)
      b = N;
    if (l < 0 || r < 0)
      assert(false);
    lazy_evaluate_node(
        k, a,
        b); // ??¨?????????????????£???????????????????????????????????§?????¬?????¨???????\\?.

    if (b <= l || r <= a) //[a,b)??¨[l,r)???????????????????????´???
      return;
    if (l <= a && b <= r) { // [l,r)???[a,b)????????¨???????????§????????´???
      dat[k].lazy = lazy_connect(dat[k].lazy, v);
      lazy_evaluate_node(
          k, a, b); //???????????¶???????????¨???????????¨?????????????????§.
      return;
    }

    int m = (a + b) / 2;
    update(l, r, v, 2 * k, a, m);
    update(l, r, v, 2 * k + 1, m, b);
    update_node(k);
  }
  // get(l,r) := [l,r)?????????????????¨?????????????????????.
  // ?????????0-indexed.
  Value1 get(int l, int r, int k = 1, int a = 0, int b = -1) {
    if (b == -1)
      b = N;

    if (l < 0 || r < 0)
      assert(false);
    lazy_evaluate_node(
        k, a,
        b); // ??¨?????????????????£???????????????????????????????????§?????¬?????¨???????\\?.

    if (b <= l || r <= a) //[a,b)??¨[l,r)???????????????????????´???
      return Zero1;

    if (l <= a && b <= r) { // [l,r)???[a,b)????????¨???????????§????????´???
      return dat[k].sum;
    }

    int m = (a + b) / 2;
    Value1 vl = get(l, r, 2 * k, a, m);
    Value1 vr = get(l, r, 2 * k + 1, m, b);
    update_node(k);
    return connect(vl, vr);
  }
};

void dfs(const vector<vector<int>> &edges, const int now, const int from,
         vector<int> &ls, vector<int> &rs, int &id) {
  ls[now] = id;
  id++;
  for (auto e : edges[now]) {
    if (e == from)
      continue;
    dfs(edges, e, now, ls, rs, id);
  }
  rs[now] = id;
}

int main() {
  cin >> aN >> Q;
  vector<vector<int>> edges(aN);
  vector<int> ls(aN), rs(aN);
  for (int i = 0; i < aN - 1; ++i) {
    int a, b;
    cin >> a >> b;
    edges[a].push_back(b);
    edges[b].push_back(a);
  }
  int id = 0;
  dfs(edges, 0, -1, ls, rs, id);

  lazy_segtree seg(aN);

  while (Q--) {
    int t, r, x, y;
    cin >> t >> r >> x >> y;
    if (t == 1) {
      Value1 v1 = seg.get(ls[r], rs[r]);
      int ans = accumulate(v1.begin() + x, v1.begin() + y + 1, 0);
      cout << ans << endl;
    } else {
      vector<short> v(11);
      iota(v.begin(), v.end(), 0);
      v[x] = y;
      seg.update(ls[r], rs[r], v);
    }
  }

  return 0;
}

Error: TLE

Fixed submission:

#include "bits/stdc++.h"
#include <unordered_map>
#include <unordered_set>
#pragma warning(disable : 4996)
using namespace std;

typedef vector<int> Value1;
typedef vector<short> Value2;
const Value1 Zero1{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1};
const Value2 Zero2{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10};

struct Node {
  Value1
      sum; //??´??°????????????.
           //?????????????????§??????????????????????????¨????????????????????????????????????.
  Value2 lazy; //?????¶???????????????????????????????????????
  Node() : sum(Zero1) { lazy = Zero2; }
};
struct lazy_segtree {
  int N;
  vector<Node> dat;
  lazy_segtree(int n) : N(1) {
    while (N < n)
      N *= 2;
    dat.resize(2 * N);
  }

  Value2 lazy_connect(const Value2 l, const Value2 r) {
    Value2 v(l);
    for (int i = 0; i < 11; ++i) {
      v[i] = r[v[i]];
    }
    return v;
  }

  void lazy_func(const int k, const int a, const int b) {
    Value1 next(11);
    for (int i = 0; i < 11; ++i) {
      next[dat[k].lazy[i]] += dat[k].sum[i];
    }
    dat[k].sum = next;
  }

  Value1 connect(const Value1 l, const Value1 r) {
    Value1 next(l);
    for (int i = 0; i < 11; ++i)
      next[i] += r[i];
    return next;
  }

  // inline??????????????¨??§???????????¨????????????!(??????)
  inline void lazy_evaluate_node(int k, int a, int b) {
    lazy_func(k, a, b);
    if (k <
        N) { // 2*k(???????????????) < 2*N (???????????°) ?????????????????§.
             // ???????????????????????????????????¬???????????¨??????.
      dat[2 * k].lazy = lazy_connect(
          dat[2 * k].lazy,
          dat[k].lazy); //?¬???????????????¬??????????????£?????????.
      dat[2 * k + 1].lazy = lazy_connect(dat[2 * k + 1].lazy, dat[k].lazy);
    }
    dat[k].lazy = Zero2;
  }

  inline void update_node(
      int k) { // k???????????¢?????????????????????????????¨?????????.
               // ????????\\????????¨????????????????????????????????????????????????????????§if???????????????.
    dat[k].sum = connect(dat[2 * k].sum, dat[2 * k + 1].sum);
  }

  // update(l,r,v) := [l,r)?????´??°??????. ?????????0-indexed.
  void update(int l, int r, const Value2 &v, int k = 1, int a = 0, int b = -1) {
    if (b == -1)
      b = N;
    if (l < 0 || r < 0)
      assert(false);
    lazy_evaluate_node(
        k, a,
        b); // ??¨?????????????????£???????????????????????????????????§?????¬?????¨???????\\?.

    if (b <= l || r <= a) //[a,b)??¨[l,r)???????????????????????´???
      return;
    if (l <= a && b <= r) { // [l,r)???[a,b)????????¨???????????§????????´???
      dat[k].lazy = lazy_connect(dat[k].lazy, v);
      lazy_evaluate_node(
          k, a, b); //???????????¶???????????¨???????????¨?????????????????§.
      return;
    }

    int m = (a + b) / 2;
    update(l, r, v, 2 * k, a, m);
    update(l, r, v, 2 * k + 1, m, b);
    update_node(k);
  }
  // get(l,r) := [l,r)?????????????????¨?????????????????????.
  // ?????????0-indexed.
  Value1 get(int l, int r, int k = 1, int a = 0, int b = -1) {
    if (b == -1)
      b = N;

    if (l < 0 || r < 0)
      assert(false);
    lazy_evaluate_node(
        k, a,
        b); // ??¨?????????????????£???????????????????????????????????§?????¬?????¨???????\\?.

    if (b <= l || r <= a) //[a,b)??¨[l,r)???????????????????????´???
      return Zero1;

    if (l <= a && b <= r) { // [l,r)???[a,b)????????¨???????????§????????´???
      return dat[k].sum;
    }

    int m = (a + b) / 2;
    Value1 vl = get(l, r, 2 * k, a, m);
    Value1 vr = get(l, r, 2 * k + 1, m, b);
    update_node(k);
    return connect(vl, vr);
  }
};

void dfs(const vector<vector<int>> &edges, const int now, const int from,
         vector<int> &ls, vector<int> &rs, int &id) {
  ls[now] = id;
  id++;
  for (auto e : edges[now]) {
    if (e == from)
      continue;
    dfs(edges, e, now, ls, rs, id);
  }
  rs[now] = id;
}

int main() {
  int N, Q;
  cin >> N >> Q;
  vector<vector<int>> edges(N);
  vector<int> ls(N), rs(N);
  for (int i = 0; i < N - 1; ++i) {
    int a, b;
    cin >> a >> b;
    edges[a].push_back(b);
    edges[b].push_back(a);
  }
  int id = 0;
  dfs(edges, 0, -1, ls, rs, id);

  lazy_segtree seg(N);
  vector<short> v{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0};
  for (int i = 0; i < N; ++i) {
    seg.update(i, i + 1, v);
  }
  while (Q--) {
    int t, r, x, y;
    cin >> t >> r >> x >> y;
    if (t == 1) {
      Value1 v1 = seg.get(ls[r], rs[r]);
      int ans = accumulate(v1.begin() + x, v1.begin() + y + 1, 0);
      cout << ans << endl;
    } else {
      vector<short> v(11);
      iota(v.begin(), v.end(), 0);
      v[x] = y;
      seg.update(ls[r], rs[r], v);
    }
  }

  return 0;
}