Add lazy segment tree template

Author: Mr-Python-in-China

.gitignore

@@ -6,3 +6,4 @@ node_modules
 library/test.exe
 library/test.d
 .yarn/
+*.gch

.vscode/settings.json

@@ -76,5 +76,9 @@
     "typeinfo": "cpp",
     "variant": "cpp",
     "text_encoding": "cpp"
-  }
-}
+  },
+  "clangd.fallbackFlags": [
+    "-std=c++14",
+    "-I${workspaceFolder}/library"
+  ]
+}

.clang-format renamed to library/.clang-format

@@ -0,0 +1,125 @@
+#include "mrpython/bits.hpp"
+
+#ifndef MP_LIBRARY_LAZY_SEGMENT_TREE_HPP
+#define MP_LIBRARY_LAZY_SEGMENT_TREE_HPP
+#include <algorithm>
+#include <functional>
+#include <vector>
+
+namespace mrpython {
+using std::size_t;
+template <typename T, typename MergeFunction, typename Lazy,
+          typename OperateFunction, typename MergeLazyFunction>
+class lazy_segment_tree {
+  std::vector<T> data;
+  std::vector<Lazy> lazy;
+  std::vector<size_t> size;
+  size_t n;
+  Lazy const lazyInit;
+  MergeFunction const mergeData;
+  OperateFunction const operate;
+  MergeLazyFunction const mergeLazy;
+  void build(void) {
+    data.reserve(2 * n - 1), size.reserve(2 * n - 1);
+    for (size_t i = n; i < 2 * n - 1; ++i) {
+      size_t d = 2 * n - 1 - i;
+      size_t l = d * 2 + 1, r = d * 2;
+      data.emplace_back(mergeData(data[2 * n - 1 - l], data[2 * n - 1 - r]));
+      size.emplace_back(size[2 * n - 1 - l] + size[2 * n - 1 - r]);
+    }
+    std::reverse(data.begin(), data.end());
+    std::reverse(size.begin(), size.end());
+  }
+  void pushdown(size_t pos) {
+    if (size[pos] == 1) {
+      lazy[pos] = lazyInit;
+      return;
+    }
+    add_tag_for_node(2 * pos + 1, lazy[pos]);
+    add_tag_for_node(2 * pos + 2, lazy[pos]);
+    lazy[pos] = lazyInit;
+  }
+  void add_tag_for_node(size_t pos, Lazy const& lazyVal) {
+    data[pos] = operate(lazyVal, data[pos], size[pos]);
+    lazy[pos] = mergeLazy(lazy[pos], lazyVal);
+  }
+  T get_impl(size_t l, size_t r, size_t pos) {
+    if (l == 0 && r == size[pos]) return data[pos];
+    pushdown(pos);
+    size_t m = size[pos * 2 + 1];
+    if (l < m && r > m)
+      return mergeData(get_impl(l, m, pos * 2 + 1),
+                       get_impl(0, r - m, pos * 2 + 2));
+    else if (l < m)
+      return get_impl(l, r, pos * 2 + 1);
+    else if (r > m)
+      return get_impl(l - m, r - m, pos * 2 + 2);
+    else
+      __builtin_unreachable();
+  }
+  void set_impl(size_t l, size_t r, Lazy const& operateVal, size_t pos) {
+    if (l == 0 && r == size[pos]) return add_tag_for_node(pos, operateVal);
+    pushdown(pos);
+    size_t m = size[pos * 2 + 1];
+    if (l < m && r > m)
+      set_impl(l, m, operateVal, pos * 2 + 1),
+          set_impl(0, r - m, operateVal, pos * 2 + 2);
+    else if (l < m)
+      set_impl(l, r, operateVal, pos * 2 + 1);
+    else if (r > m)
+      set_impl(l - m, r - m, operateVal, pos * 2 + 2);
+    else
+      __builtin_unreachable();
+    data[pos] = mergeData(data[pos * 2 + 1], data[pos * 2 + 2]);
+  }
+
+ public:
+  template <typename InputIterator>
+  lazy_segment_tree(InputIterator first, InputIterator last,
+                    Lazy const& lazyInitVal,
+                    MergeFunction const& mergeDataFun = MergeFunction(),
+                    OperateFunction const& OperateFun = OperateFunction(),
+                    MergeLazyFunction const& mergeTagFun = MergeLazyFunction())
+      : data(first, last),
+        lazy(2 * data.size() - 1, lazyInitVal),
+        size(data.size(), 1),
+        n(data.size()),
+        lazyInit(lazyInitVal),
+        mergeData(mergeDataFun),
+        operate(OperateFun),
+        mergeLazy(mergeTagFun) {
+    rotate(data.begin(), data.begin() + (2 * n - 1) - (highbit(2 * n - 1) - 1),
+           data.end());
+    reverse(data.begin(), data.end());
+    build();
+  }
+  lazy_segment_tree(size_t len, T const& init, Lazy const& lazyInitVal,
+                    MergeFunction const& mergeDataFun = MergeFunction(),
+                    OperateFunction const& OperateFun = OperateFunction(),
+                    MergeLazyFunction const& mergeTagFun = MergeLazyFunction())
+      : data(len, init),
+        lazy(2 * len - 1, lazyInitVal),
+        size(len, 1),
+        n(len),
+        lazyInit(lazyInitVal),
+        mergeData(mergeDataFun),
+        mergeLazy(mergeTagFun) {
+    build();
+  }
+  T get(size_t l, size_t r) { return get_impl(l, r, 0); }
+  void set(size_t l, size_t r, Lazy const& operateVal) {
+    set_impl(l, r, operateVal, 0);
+  }
+};
+template <typename T> struct lazy_segment_tree_add_add_operate_function {
+  T operator()(T const& lazy, T const& data, size_t size) const {
+    return data + lazy * size;
+  }
+};
+template <typename T>
+using lazy_segment_tree_add_add =
+    lazy_segment_tree<T, std::plus<T>, T,
+                      lazy_segment_tree_add_add_operate_function<T>,
+                      std::plus<T>>;
+}  // namespace mrpython
+#endif

library/mrpython/lazy_segment_tree.hpp

@@ -4,16 +4,14 @@
 #ifndef MP_LIBRARY_SEGMENT_TREE_HPP
 #define MP_LIBRARY_SEGMENT_TREE_HPP
 #include <algorithm>
-#include <iterator>
-#include <utility>
 #include <vector>
 
 namespace mrpython {
 using std::size_t;
 /**
  * 一颗 无标记线段树 模板。
  * @tparam T 元素类型
- * @tparam MergeFunction 合并运算函数，对于函数 `F` 应满足结合律
+ * @tparam MergeFunction 合并运算函数
  */
 template <typename T, typename MergeFunction> class typical_segment_tree {
   std::vector<T> data;
@@ -31,10 +29,10 @@ template <typename T, typename MergeFunction> class typical_segment_tree {
     std::reverse(data.begin(), data.end());
     std::reverse(size.begin(), size.end());
   }
-  template <typename Fun>
-  void set_impl(size_t c, Fun const& operate, size_t pos) {
+  template <typename Operate>
+  void set_impl(size_t c, Operate const& operate, size_t pos) {
     if (size[pos] == 1) {
-      data[pos] = operate((typename std::vector<T>::const_reference)data[pos]);
+      data[pos] = operate((T const&)data[pos]);
       return;
     }
     size_t m = size[pos * 2 + 1];
@@ -50,15 +48,18 @@ template <typename T, typename MergeFunction> class typical_segment_tree {
     if (l < m && r > m)
       return merge(get_impl(l, m, pos * 2 + 1),
                    get_impl(0, r - m, pos * 2 + 2));
-    if (l < m) return get_impl(l, r, pos * 2 + 1);
-    if (r > m) return get_impl(l - m, r - m, pos * 2 + 2);
-    throw;
+    else if (l < m)
+      return get_impl(l, r, pos * 2 + 1);
+    else if (r > m)
+      return get_impl(l - m, r - m, pos * 2 + 2);
+    else
+      __builtin_unreachable();
   }
 
  public:
   template <typename InputIterator>
   typical_segment_tree(InputIterator first, InputIterator last,
-                       MergeFunction mergeFun = MergeFunction())
+                       MergeFunction const& mergeFun = MergeFunction())
       : data(first, last),
         size(data.size(), 1),
         n(data.size()),
@@ -69,7 +70,7 @@ template <typename T, typename MergeFunction> class typical_segment_tree {
     build();
   }
   typical_segment_tree(size_t len, T const& init,
-                       MergeFunction mergeFun = MergeFunction())
+                       MergeFunction const& mergeFun = MergeFunction())
       : data(len, init), size(len, 1), n(len), merge(mergeFun) {
     build();
   }

library/mrpython/typical_segment_tree.hpp

@@ -1,2 +1,3 @@
+#include "lazy_segment_tree/all.hpp"
 #include "sparse_table/all.hpp"
 #include "typical_segment_tree/all.hpp"

library/test/all.hpp

@@ -0,0 +1 @@
+#include "lazy_segment_tree_add_add.cpp"

library/test/lazy_segment_tree/all.hpp

@@ -0,0 +1,30 @@
+#include <gtest/gtest.h>
+
+#include <mrpython/lazy_segment_tree.hpp>
+#include <random>
+
+TEST(lazy_segment_tree, add) {
+  std::mt19937_64 gen(std::random_device{}());
+  std::size_t n = std::uniform_int_distribution<std::size_t>{1, 5000}(gen),
+              q = std::uniform_int_distribution<std::size_t>{1, 5000}(gen);
+  std::uniform_int_distribution<unsigned> val_dist(
+      std::numeric_limits<unsigned>::min(),
+      std::numeric_limits<unsigned>::max()),
+      size_dist(0, n - 1), operator_dist(0, 1);
+  std::vector<unsigned> a(n);
+  std::generate(a.begin(), a.end(), [&] { return val_dist(gen); });
+  mrpython::lazy_segment_tree_add_add<unsigned> tree(a.begin(), a.end(), 0);
+  while (q--) {
+    std::size_t l = size_dist(gen), r = size_dist(gen);
+    if (l > r) std::swap(l, r);
+    assert(l < r + 1);
+    if (operator_dist(gen)) {
+      int ans = std::accumulate(a.begin() + l, a.begin() + r + 1, (unsigned)0);
+      EXPECT_EQ(tree.get(l, r + 1), ans);
+    } else {
+      unsigned value = val_dist(gen);
+      for (std::size_t i = l; i < r; ++i) a[i] += value;
+      tree.set(l, r, value);
+    }
+  }
+}