verify/Datastructure_treap.test.cpp

• View this file on GitHub
• Last update: 2025-01-10 00:00:54+09:00
• Problem: https://judge.yosupo.jp/problem/dynamic_sequence_range_affine_range_sum

Depends on

• treap (Datastructure/treap.hpp)
• modint (Utility/modint.hpp)
• verify用テンプレート (Utility/template.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/dynamic_sequence_range_affine_range_sum"
#include "../Utility/template.hpp"
#include "../Utility/modint.hpp"
#include "../Datastructure/treap.hpp"
using mint = modint<998244353>;

struct S {
     mint s; int sz;
};

S op(S l, S r) {
    return S{l.s + r.s, l.sz + r.sz};
}

S e() {
    return S{mint(0), 0};
}

struct F {
    mint a, b;
};

S mapping(F f, S s) {
    S res;
    res.s = f.a * s.s + f.b * s.sz;
    res.sz = s.sz;
    return res;
}

F composition(F l, F r) {
    return F{l.a * r.a,  l.a * r.b + l.b};
}

F id() {
    return F{1, 0};
}


int main() {
    treap<S, op, e, F, mapping, composition, id> tr;
    ll N, Q;
    cin >> N >> Q;
    for(int i = 0; i <= N-1; i++) {
        ll t;
        cin >> t;
        mint a = t;
        tr.push_back(S{a,1});
    }

    while(Q--) {
        assert(tr.size() == N);
        int type;
        cin >> type;
        if(type == 0) {
            ll i, x;
            cin >> i >> x;
            tr.insert(i, S{x, 1});
            N++;
        }
        else if(type==1) {
            int i;
            cin >> i;
            tr.erase(i);
            N--;
        }
        else if(type==2) {
            ll l, r;
            cin >> l >> r;
            tr.reverse(l,r);
        }
        else if(type==3) {
            ll l, r;
            mint b, c;
            cin >> l >> r >> b >> c;
            tr.apply(l, r, F{b, c});
        }
        else {
            ll l, r;
            cin >> l >> r;
            if(l==0 && r == tr.size()) {
                cout << tr.all_prod().s.x << " ";
                continue;
            }
            cout << tr.prod(l, r).s.x << " ";
        }
    }
}

#line 1 "verify/Datastructure_treap.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/dynamic_sequence_range_affine_range_sum"
#line 1 "Utility/template.hpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define rep(i, s, t) for (ll i = s; i < (ll)(t); i++)
#define rrep(i, s, t) for (ll i = (ll)(t) - 1; i >= (ll)(s); i--)
#define all(x) begin(x), end(x)

#define TT template <typename T>
TT using vec = vector<T>;
template <class T1, class T2> bool chmin(T1 &x, T2 y) {
    return x > y ? (x = y, true) : false;
}
template <class T1, class T2> bool chmax(T1 &x, T2 y) {
    return x < y ? (x = y, true) : false;
}
struct io_setup {
    io_setup() {
        ios::sync_with_stdio(false);
        std::cin.tie(nullptr);
        cout << fixed << setprecision(15);
    }
} io_setup;

/*
@brief verify用テンプレート
*/
#line 1 "Utility/modint.hpp"

// 動的mod : template<int mod> を消して、上の方で変数modを宣言
template <uint32_t mod> struct modint {
    using mm = modint;
    uint32_t x;
    modint() : x(0) {}
    TT modint(T a = 0) : x((ll(a) % mod + mod)) {
        if (x >= mod) x -= mod;
    }

    friend mm operator+(mm a, mm b) {
        a.x += b.x;
        if (a.x >= mod) a.x -= mod;
        return a;
    }
    friend mm operator-(mm a, mm b) {
        a.x -= b.x;
        if (a.x >= mod) a.x += mod;
        return a;
    }

    mm operator-() const { return mod - x; }

    //+と-だけで十分な場合、以下は省略して良いです。

    friend mm operator*(mm a, mm b) { return (uint64_t)(a.x) * b.x; }
    friend mm operator/(mm a, mm b) { return a * b.inv(); }
    friend mm &operator+=(mm &a, mm b) { return a = a + b; }
    friend mm &operator-=(mm &a, mm b) { return a = a - b; }
    friend mm &operator*=(mm &a, mm b) { return a = a * b; }
    friend mm &operator/=(mm &a, mm b) { return a = a * b.inv(); }

    mm inv() const {
        assert(x != 0);
        return pow(mod - 2);
    }
    mm pow(ll y) const {
        mm res = 1;
        mm v = *this;
        while (y) {
            if (y & 1) res *= v;
            v *= v;
            y /= 2;
        }
        return res;
    }

    friend istream &operator>>(istream &is, mm &a) {
        ll t;
        cin >> t;
        a = mm(t);
        return is;
    }

    friend ostream &operator<<(ostream &os, mm a) { return os << a.x; }

    bool operator==(mm a) { return x == a.x; }
    bool operator!=(mm a) { return x != a.x; }

    bool operator<(const mm &a) const { return x < a.x; }
};
using modint998244353 = modint<998244353>;
using modint1000000007 = modint<1'000'000'007>;
/*
@brief modint
*/
#line 1 "Datastructure/treap.hpp"
class xorshift {
    uint64_t x;
    public:
        xorshift() {
            mt19937 rnd(chrono::steady_clock::now().time_since_epoch().count());
            x = rnd();
            for (int i = 0; i < 100; i++) {
                random();
            }
        }
        uint64_t random() {
            x = x ^ (x << 7);
            return x = x ^ (x >> 9);
    }
};


template<class S, S (*op)(S, S), S (*e)(), class F, S (*mapping)(F, S), F (*composition)(F, F), F (*id)()>
struct treap {
    xorshift rnd;
    int sz = 0; 
    private:
        struct node_t{
            node_t* lch;
            node_t* rch;
            int pri, cnt;
            S val, acc;
            F lazy;
            bool rev;
            bool have_e;
 
            node_t(S v, int p) : val(v), pri(p), acc(v) , lch(nullptr), rch(nullptr), rev(false), cnt(1) {
                lch = rch = nullptr;
                rev = false;
                have_e = false;
                lazy = id();
            }
        };

        using np = node_t*;

        np root = nullptr;

        long long count(np t) {return !t ? 0 : t -> cnt;}
        
        S acc(np t) {return !t ? e() : t -> acc; }

        np pushup(np t) {
            if(t) {
                t -> cnt = count(t -> lch) + count(t -> rch) + 1;
                t -> acc = op(op(acc(t -> lch), t -> val),  acc(t -> rch));
            }
            return t;
        }
 
        void pushdown(node_t *t) {
            if(t && t -> rev) {
                swap(t -> lch, t -> rch);
                if(t -> lch) t -> lch -> rev ^= 1;
                if(t -> rch) t -> rch -> rev ^= 1;
                t -> rev = false;
            }

            if(t && t -> have_e) {
                if(t -> lch) {
                    t -> lch-> lazy = composition(t -> lazy, t -> lch -> lazy);
                    t -> lch -> acc = mapping(t -> lazy, t -> lch -> acc);
                    t -> lch -> have_e = true;
                }

                if(t -> rch) {
                    t -> rch -> lazy = composition(t -> lazy, t -> rch -> lazy);
                    t -> rch -> acc = mapping(t -> lazy, t -> rch -> acc);
                    t -> rch -> have_e = true;
                }
                t -> val = mapping(t -> lazy, t -> val);
                t -> lazy = id();
                t -> have_e = false;
            }
            pushup(t);
        }
 
        void merge(np& t, np l, np r) {
            pushdown(l), pushdown(r);
            if(!l || !r) t =  !l ? r : l;
            else if(l -> pri > r -> pri) {
                merge(l -> rch, l -> rch, r);
                t = l;
            } else {
               merge(r -> lch, l,r -> lch);
               t = r;
            }
            pushup(t);
        }

        void split(np t, int k, np& tl, np& tr) {// [0, k) [k, n)
            if(!t) {
                tl = nullptr, tr = nullptr;
                return;
            }
            pushdown(t);
 
            if(k <= count(t -> lch)) {
                split(t -> lch, k, tl, t -> lch);
                tr = t;
            }else {
                split(t -> rch, k - count(t -> lch) - 1, t -> rch, tr);
                tl = t;
            }
            pushup(t);
        }

        
        void dump__(np t, vector<S>& res) {
            if(!t) return;
            pushdown(t);
            dump__(t -> lch, res);
            res.push_back(t -> val);
            dump__(t -> rch, res);
        }


    public:
        void insert(int p, S val) {
            assert(0 <= p && p <= size());
            np nw = new node_t(val, rnd.random());
            np tl; np tr;
            split(root, p, tl, tr);
            merge(tl, tl, nw);
            merge(root, tl, tr);
            sz++;
        }

        void push_back(S val) {insert(size(), val);}
 
        void erase(int p) {
            assert(0 <= p && p < size());
            np tl; np tm; np tr;
            split(root, p+1, tl, tm);
            split(tl, p, tl, tr);
            merge(root, tl, tm);
            sz--;
        }

        void pop_back() {
            assert(size()>0);
            erase(size()-1);
        }


        S prod(int l, int r) {
            if(l >= r) return e();
            assert(0 <= l && r <= size());
            np tl; np tm; np tr;
            split(root, l, tl, tm);
            split(tm, r-l, tm, tr);
            S res = acc(tm);
            merge(tm, tm, tr);
            merge(root, tl, tm);
            return res;
        }

        S all_prod() {
            assert(size() > 0);
            pushdown(root);
            pushup(root);
            return root -> acc;
        }

        S get(int p) {
            assert(0 <= p && p < size());
            return prod(p, p+1);
        }

        void apply(int p, F f) {apply(p, p+1, f);}

        void apply(int l, int r, F f) {
            if(l >= r) return;
            assert(0 <= l && r <= size());
            np tl; np tm; np tr;
            split(root, l, tl, tm);
            split(tm, r - l, tm, tr);
            tm -> have_e = true;
            tm -> lazy = composition(tm -> lazy, f);
            tm -> acc = mapping(f, tm -> acc);
            merge(tm, tm, tr);
            merge(root, tl, tm);
        }


        void reverse(int l, int r) {//[l, r)をreverse
            if(l >= r) return;
            assert(0 <= l && r <= size());
            np tl; np tm; np tr;
            split(root, l, tl, tm);
            split(tm, r - l, tm, tr);
            tm -> rev ^= 1;
            merge(tm, tm, tr);
            merge(root, tl, tm);
        }
 
        void rotate(int l, int m, int r) {//[l, r) を mが先頭に来る様にreverse
            if(l >= r) return;
            assert(l <= m && m < r);
            assert(0 <= l && r <= size());
            reverse(l, r);
            reverse(l, l + r - m);
            reverse(l + r - m, r);
        }
      
        vector<S> dump() {
            vector<S> res;
            dump__(root, res);
            return res;
        }

        int size() {
            return sz;
        }
 
};
/*
@brief treap
@docs doc/treap.md
*/
#line 5 "verify/Datastructure_treap.test.cpp"
using mint = modint<998244353>;

struct S {
     mint s; int sz;
};

S op(S l, S r) {
    return S{l.s + r.s, l.sz + r.sz};
}

S e() {
    return S{mint(0), 0};
}

struct F {
    mint a, b;
};

S mapping(F f, S s) {
    S res;
    res.s = f.a * s.s + f.b * s.sz;
    res.sz = s.sz;
    return res;
}

F composition(F l, F r) {
    return F{l.a * r.a,  l.a * r.b + l.b};
}

F id() {
    return F{1, 0};
}


int main() {
    treap<S, op, e, F, mapping, composition, id> tr;
    ll N, Q;
    cin >> N >> Q;
    for(int i = 0; i <= N-1; i++) {
        ll t;
        cin >> t;
        mint a = t;
        tr.push_back(S{a,1});
    }

    while(Q--) {
        assert(tr.size() == N);
        int type;
        cin >> type;
        if(type == 0) {
            ll i, x;
            cin >> i >> x;
            tr.insert(i, S{x, 1});
            N++;
        }
        else if(type==1) {
            int i;
            cin >> i;
            tr.erase(i);
            N--;
        }
        else if(type==2) {
            ll l, r;
            cin >> l >> r;
            tr.reverse(l,r);
        }
        else if(type==3) {
            ll l, r;
            mint b, c;
            cin >> l >> r >> b >> c;
            tr.apply(l, r, F{b, c});
        }
        else {
            ll l, r;
            cin >> l >> r;
            if(l==0 && r == tr.size()) {
                cout << tr.all_prod().s.x << " ";
                continue;
            }
            cout << tr.prod(l, r).s.x << " ";
        }
    }
}

Back to top page