make segtree to allow more widely functions

Author: yosupo06

atcoder/lazysegtree.hpp

@@ -10,6 +10,31 @@
 
 namespace atcoder {
 
+#if __cplusplus >= 201703L
+
+template <class S,
+          auto op,
+          auto e,
+          class F,
+          auto mapping,
+          auto composition,
+          auto id>
+struct lazy_segtree {
+    static_assert(std::is_convertible_v<decltype(op), std::function<S(S, S)>>,
+                  "op must work as S(S, S)");
+    static_assert(std::is_convertible_v<decltype(e), std::function<S()>>,
+                  "e must work as S()");
+    static_assert(
+        std::is_convertible_v<decltype(mapping), std::function<S(F, S)>>,
+        "mapping must work as F(F, S)");
+    static_assert(
+        std::is_convertible_v<decltype(composition), std::function<F(F, F)>>,
+        "compostiion must work as F(F, F)");
+    static_assert(std::is_convertible_v<decltype(id), std::function<F()>>,
+                  "id must work as F()");
+
+#else
+
 template <class S,
           S (*op)(S, S),
           S (*e)(),
@@ -18,6 +43,9 @@ template <class S,
           F (*composition)(F, F),
           F (*id)()>
 struct lazy_segtree {
+
+#endif
+
   public:
     lazy_segtree() : lazy_segtree(0) {}
     explicit lazy_segtree(int n) : lazy_segtree(std::vector<S>(n, e())) {}

atcoder/segtree.hpp

@@ -6,10 +6,22 @@
 #include <vector>
 
 #include "atcoder/internal_bit"
-
 namespace atcoder {
 
+#if __cplusplus >= 201703L
+
+template <class S, auto op, auto e> struct segtree {
+    static_assert(std::is_convertible_v<decltype(op), std::function<S(S, S)>>,
+                  "op must work as S(S, S)");
+    static_assert(std::is_convertible_v<decltype(e), std::function<S()>>,
+                  "e must work as S()");
+
+#else
+
 template <class S, S (*op)(S, S), S (*e)()> struct segtree {
+
+#endif
+
   public:
     segtree() : segtree(0) {}
     explicit segtree(int n) : segtree(std::vector<S>(n, e())) {}

test/unittest/lazysegtree_test.cpp

@@ -85,3 +85,33 @@ TEST(LazySegtreeTest, Usage) {
     ASSERT_EQ(-5, seg.prod(2, 3));
     ASSERT_EQ(0, seg.prod(2, 4));
 }
+
+#if __cplusplus >= 201703L
+
+int op_const(const int& a, const int& b) { return std::max(a, b); }
+
+struct const_starry {
+    static int op_ss(const int& a, const int& b) { return std::max(a, b); }
+    static int op_ts(const int& a, const int& b) { return a + b; }
+    static int op_tt(const int& a, const int& b) { return a + b; }
+};
+
+TEST(SegtreeTest, ConstFunc) {
+    lazy_segtree<int, const_starry::op_ss, starry::e_s, int,
+                 const_starry::op_ts, const_starry::op_tt, starry::e_t>
+        seg(10);
+}
+
+#endif
+
+#if __cplusplus >= 202002L
+
+TEST(LazySegtreeTest, LambdaFunc) {
+    lazy_segtree<int, [](int a, int b) { return std::max(a, b); },
+                 []() { return -1'000'000'000; }, int,
+                 [](int a, int b) { return a + b; },
+                 [](int a, int b) { return a + b; }, []() { return 0; }>
+        seg(10);
+}
+
+#endif

test/unittest/segtree_test.cpp

@@ -144,3 +144,28 @@ TEST(SegtreeTest, Assign) {
     seg seg0;
     seg0 = seg(10);
 }
+
+#if __cplusplus >= 201703L
+
+std::string op_const(const std::string& a, const std::string& b) {
+    assert(a == "$" || b == "$" || a <= b);
+    if (a == "$") return b;
+    if (b == "$") return a;
+    return a + b;
+}
+
+TEST(SegtreeTest, ConstFunc) {
+    segtree<std::string, op_const, e> s1(10);
+}
+
+#endif
+
+#if __cplusplus >= 202002L
+
+TEST(SegtreeTest, LambdaFunc) {
+    segtree<std::string, [](std::string a, std::string b) {
+        return a + b;
+    }, []() { return ""; }> s1(10);
+}
+
+#endif