Advertisement
naquo

LazyST

Nov 20th, 2019
99
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 3.04 KB | None | 0 0
  1. enum Spell { NONE, INVERT, BUCCANEERIZE, BARBARIZE };
  2.  
  3. using DATA = int;
  4. using LAZYDATA = Spell;
  5. const DATA NEUT = 0;
  6. const LAZYDATA LAZYNEUT = NONE;
  7. const auto OP = std::plus<DATA>();
  8. const auto APPLY = [](const DATA& a, const LAZYDATA& spell, int len) -> DATA {
  9.     if (spell == NONE) return a;
  10.     if (spell == INVERT) return len - a;
  11.     if (spell == BUCCANEERIZE) return len;
  12.     if (spell == BARBARIZE) return 0;
  13.     assert(not "Unknown spell");
  14. };
  15. const auto COMBINE = [](const LAZYDATA& a, const LAZYDATA& b) -> LAZYDATA {
  16.     if (a == NONE) return b;
  17.     if (b == NONE) return a;
  18.     if (b == BUCCANEERIZE || b == BARBARIZE) return b;
  19.     // b must be INVERT
  20.     assert(b == INVERT);
  21.     if (a == BUCCANEERIZE) return BARBARIZE;
  22.     if (a == BARBARIZE) return BUCCANEERIZE;
  23.     if (a == INVERT) return NONE;
  24.     assert(not "Unknown spell combination");
  25. };
  26.  
  27. struct LazyST {
  28.     std::vector<DATA> st;
  29.     std::vector<LAZYDATA> lazy;
  30.  
  31.     // internal functions
  32.     inline int parent(int i) const { return i / 2; }
  33.     inline int lchild(int i) const { return 2 * i; }
  34.     inline int rchild(int i) const { return 2 * i + 1; }
  35.     inline int leaf_offset() const { return st.size() / 2; }
  36.     int next_pow_2(int i) const {
  37.         int r = 1;
  38.         while (r < i)
  39.             r *= 2;
  40.         return r;
  41.     }
  42.  
  43.     void apply(int idx, int len) {
  44.         if (lazy[idx] == LAZYNEUT) return;
  45.         st[idx] = APPLY(st[idx], lazy[idx], len);
  46.  
  47.         if (len > 1) {
  48.             lazy[lchild(idx)] =
  49.                 COMBINE(lazy[lchild(idx)], lazy[idx]);
  50.             lazy[rchild(idx)] =
  51.                 COMBINE(lazy[rchild(idx)], lazy[idx]);
  52.         }
  53.  
  54.         lazy[idx] = LAZYNEUT;
  55.     }
  56.  
  57.     DATA query(int l, int r, int idx, int i, int j) {
  58.         // return neutral value if disjoint
  59.         if (r <= i || l >= j) return NEUT;
  60.         // lower the updates one level
  61.         apply(idx, j - i);
  62.         // return node if fully contained in range
  63.         if (l <= i && j <= r) return st[idx];
  64.         // combine children if partially contained
  65.         int m = (i + j) / 2;
  66.         return OP(query(l, r, lchild(idx), i, m),
  67.                   query(l, r, rchild(idx), m, j));
  68.     }
  69.  
  70.     // API functions
  71.     // query operation - apply query from top node
  72.     DATA query(int l, int r) {
  73.         return query(l, r, 1, 0, st.size() / 2);
  74.     }
  75.  
  76.     void set(int l, int r, int idx, int i, int j, const LAZYDATA& v) {
  77.         apply(idx, j - i);
  78.         if (r <= i || l >= j) return;
  79.         if (l <= i && j <= r) {
  80.             lazy[idx] = COMBINE(lazy[idx], v);
  81.             apply(idx, j - i);
  82.             return;
  83.         }
  84.         int m = (i + j) / 2;
  85.         set(l, r, lchild(idx), i, m, v);
  86.         set(l, r, rchild(idx), m, j, v);
  87.         st[idx] = OP(st[lchild(idx)], st[rchild(idx)]);
  88.     }
  89.  
  90.     void set(int l, int r, const LAZYDATA& v) {
  91.         set(l, r, 1, 0, st.size() / 2, v);
  92.     }
  93.  
  94.     // segment tree factory
  95.     template<class Iter>
  96.     LazyST& init(Iter begin, Iter end) {
  97.         size_t size = next_pow_2(std::distance(begin, end));
  98.         // allocate space for tree
  99.         st.assign(2 * size, NEUT);
  100.         lazy.assign(2 * size, LAZYNEUT);
  101.         // load values into tree
  102.         std::copy(begin, end, st.begin() + size);
  103.         // compute upper levels
  104.         while (size = parent(size))
  105.             for (size_t i = size; i < 2 * size; i++)
  106.                 st[i] = OP(st[lchild(i)], st[rchild(i)]);
  107.  
  108.         return *this;
  109.     }
  110. };
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement