Support any operations in Fenwick

Author: adamant-pwn

cp-algo/structures/fenwick.hpp

@@ -3,56 +3,92 @@
 #include <cassert>
 #include <vector>
 namespace cp_algo::structures {
-    template<typename T, typename Container = std::vector<T>>
+    template <typename Op>
+    struct inverse_op {};
+
+    template <typename T>
+    struct inverse_op<std::plus<T>> {
+        static T apply(T const& a, T const& b) {
+            return a - b;
+        }
+    };
+
+    template <typename T>
+    struct inverse_op<std::multiplies<T>> {
+        static T apply(T const& a, T const& b) {
+            return a / b;
+        }
+    };
+
+    template<typename T, std::ranges::range Container = std::vector<T>, typename Op = std::plus<T>>
     struct fenwick {
+        Op op;
         size_t n;
         Container data;
 
-        fenwick(auto &&range) {
-            assign(move(range));
+        fenwick(auto &&range, Op &&op = Op{}): op(std::move(op)) {
+            assign(std::move(range));
         }
-        void to_prefix_sums() {
+        void to_prefix_folds() {
             for(size_t i = 1; i < n; i++) {
                 if(i + (i & -i) <= n) {
-                    data[i + (i & -i)] += data[i];
+                    data[i + (i & -i)] = op(data[i + (i & -i)], data[i]);
                 }
             }
         }
         void assign(auto &&range) {
             n = size(range) - 1;
-            data = move(range);
-            to_prefix_sums();
+            data = std::move(range);
+            to_prefix_folds();
         }
-        void add(size_t x, T const& v) {
+        void update(size_t x, T const& v) {
             for(++x; x <= n; x += x & -x) {
-                data[x] += v;
+                data[x] = op(data[x], v);
             }
         }
-        // sum of [0, r)
-        T prefix_sum(size_t r) const {
+        // fold of [0, r)
+        T prefix_fold(size_t r) const {
             assert(r <= n);
-            T res = 0;
+            T res = {};
             for(; r; r -= r & -r) {
-                res += data[r];
+                res = op(res, data[r]);
             }
             return res;
         }
-        // sum of [l, r)
-        T range_sum(size_t l, size_t r) const {
-            return prefix_sum(r) - prefix_sum(l);
+        // fold of [l, r)
+        T range_fold(size_t l, size_t r) const {
+            return inverse_op<Op>::apply(prefix_fold(r), prefix_fold(l));
         }
-        // Last x s.t. k = prefix_sum(x) + r for r > 0
-        // Assumes data[x] >= 0 for all x, returns [x, r]
+        // Last x s.t. prefix_fold(x) <= k
+        // Assumes prefix_fold is monotonic
+        // returns [x, prefix_fold(x)]
         auto prefix_lower_bound(T k) const {
             int x = 0;
+            T pref = {};
             for(size_t i = std::bit_floor(n); i; i /= 2) {
-                if(x + i <= n && data[x + i] < k) {
-                    k -= data[x + i];
+                if(x + i <= n && op(pref, data[x + i]) <= k) {
+                    pref = op(pref, data[x + i]);
                     x += i;
                 }
             }
-            return std::pair{x, k};
+            return std::pair{x, pref};
         }
     };
+
+    template<std::ranges::range Container, typename Op>
+    fenwick(Container&&, Op&&) -> fenwick<std::ranges::range_value_t<Container>, Container, Op>;
+    template<std::ranges::range Container>
+    fenwick(Container&&) -> fenwick<std::ranges::range_value_t<Container>, Container>;
+
+    auto maxer = [](auto const& a, auto const& b) {
+        return std::max(a, b);
+    };
+    template<typename T, std::ranges::range Container = std::vector<T>>
+    struct fenwick_max: fenwick<T, Container, decltype(maxer)> {
+        using fenwick<T, Container, decltype(maxer)>::fenwick;
+    };
+    template<std::ranges::range Container>
+    fenwick_max(Container&&) -> fenwick_max<std::ranges::range_value_t<Container>, Container>;
+
 }
 #endif // CP_ALGO_STRUCTURES_FENWICK_HPP

cp-algo/structures/fenwick_set.hpp

@@ -22,29 +22,29 @@ namespace cp_algo::structures {
                     bits.flip(x);
                 }
             }
-            Base::to_prefix_sums();
+            Base::to_prefix_folds();
         }
         void insert(size_t x) {
             if(bits.test(x)) return;
-            Base::add(x / word, 1);
+            Base::update(x / word, 1);
             bits.flip(x);
             sz++;
         }
         void erase(size_t x) {
             if(!bits.test(x)) return;
-            Base::add(x / word, -1);
+            Base::update(x / word, -1);
             bits.flip(x);
             sz--;
         }
         size_t order_of_key(size_t x) const {
-            return Base::prefix_sum(x / word) + order_of_bit(bits.word(x / word), x % word);
+            return Base::prefix_fold(x / word) + order_of_bit(bits.word(x / word), x % word);
         }
         size_t find_by_order(size_t order) const {
             if(order >= sz) {
                 return -1;
             }
-            auto [x, remainder] = Base::prefix_lower_bound(order + 1);
-            return x * word + kth_set_bit(bits.word(x), remainder - 1);
+            auto [x, pref] = Base::prefix_lower_bound(order);
+            return x * word + kth_set_bit(bits.word(x), order - pref);
         }
         size_t lower_bound(size_t x) const {
             if(bits.test(x)) {return x;}

cp-algo/util/compress_coords.hpp

@@ -4,7 +4,7 @@
 #include <vector>
 namespace cp_algo {
     // coords is a range of reference_wrapper<T>
-    auto compress_coords(auto &coords) {
+    auto compress_coords(auto &&coords) {
         using T = std::decay_t<std::unwrap_reference_t<
             std::ranges::range_value_t<decltype(coords)>
         >>;

verify/structures/fenwick/lis.test.cpp

@@ -0,0 +1,53 @@
+// @brief Longest Increasing Subsequence
+#define PROBLEM "https://judge.yosupo.jp/problem/longest_increasing_subsequence"
+#pragma GCC optimize("Ofast,unroll-loops")
+#include "cp-algo/structures/fenwick.hpp"
+#include "cp-algo/util/compress_coords.hpp"
+#include <bits/stdc++.h>
+
+using namespace std;
+
+void solve() {
+    int n;
+    cin >> n;
+    vector<reference_wrapper<int>> coords;
+    vector<int> a(n);
+    for(auto &it: a) {
+        cin >> it;
+        coords.push_back(ref(it));
+    }
+    auto b = cp_algo::compress_coords(coords);
+    struct longest {
+        int len = 0, last = 0;
+        auto operator <=>(longest const&) const = default;
+    };
+    cp_algo::structures::fenwick_max me(vector<longest>(n + 1));
+    vector<int> pre(n);
+    for(auto [i, it]: a | views::enumerate) {
+        auto [len, last] = me.prefix_fold(it);
+        me.update(it, {len + 1, (int)i});
+        pre[i] = last;
+    }
+    auto [len, last] = me.prefix_fold(n);
+    cout << len << '\n';
+    vector<int> ans(len);
+    while(len--) {
+        ans[len] = last;
+        last = pre[last];
+    }
+    for(auto it: ans) {
+        cout << it << ' ';
+    }
+}
+
+signed main() {
+    //freopen("input.txt", "r", stdin);
+    ios::sync_with_stdio(0);
+    cin.tie(0);
+    int t; 
+    t = 1;// cin >> t;
+    while(t--) {
+        solve();
+    }
+}
+ 

verify/structures/fenwick/point_add_range_sum.test.cpp

@@ -11,14 +11,14 @@ void solve() {
     cin >> n >> q;
     vector<int64_t> a(n + 1);
     for(auto &it: a | views::drop(1)) {cin >> it;}
-    cp_algo::structures::fenwick<int64_t> me(move(a));
+    cp_algo::structures::fenwick me(move(a));
     for(int i = 0; i < q; i++) {
         int t, x, y;
         cin >> t >> x >> y;
         if(t == 0) {
             me.add(x, y);
         } else {
-            cout << me.range_sum(x, y) << '\n';
+            cout << me.range_fold(x, y) << '\n';
         }
     }
 }