verify/yosupo/data-structure/point_set_tree_path_composite_sum.test.cpp

View this file on GitHub
Last update: 2024-11-28 20:30:35+07:00
Problem: https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum"
#include "template.hpp"
#include "graph/graph-base.hpp"
#include "tree/hld.hpp"
#include "tree/static-top-tree-rerooting-dp.hpp"
#include "modular-arithmetic/montgomery-modint.hpp"

using mint = mint998;

int n;
vector<int> id;
vector<mint> a,b,c;

struct TreeDP{
    struct Path{
        mint a,b,cnt,ans;
        static Path unit(){
            return {1,0,0,0};
        }
    };
    struct Point{
        mint cnt,ans;
        static Point unit(){
            return {0,0};
        }
    };
    static Path compress(const Path &l,const Path &r){
        Path res;
        res.a=l.a*r.a;
        res.b=l.a*r.b+l.b;
        res.cnt=l.cnt+r.cnt;
        res.ans=l.ans+l.a*r.ans+l.b*r.cnt;
        return res;
    }
    static Point rake(const Point &l,const Point &r){
        return {l.cnt+r.cnt,l.ans+r.ans};
    }
    static Point add_edge(const Path &p){
        return {p.cnt,p.ans};
    }
    static Path add_vertex(const Point &p,int u){
        Path res{b[u],c[u],p.cnt+(u<n),p.ans+a[u]};
        res.ans=res.a*res.ans+res.b*res.cnt;
        return res;
    }
    static Path vertex(int u){
        Path res{b[u],c[u],u<n,a[u]};
        res.ans=res.a*res.ans+res.b*res.cnt;
        return res;
    }
};

int main(){
    cin.tie(nullptr)->sync_with_stdio(false);
    int q;
    cin >> n >> q;
    a.resize(2*n);
    b.resize(2*n-1,mint(1));
    c.resize(2*n-1);
    for(int i=0;i<n;i++){
        cin >> a[i];
    }
    Graph<void,false> g(2*n-1);
    for(int i=0;i<n-1;i++){
        int u,v;
        cin >> u >> v >> b[i+n] >> c[i+n];
        g.add_edge(u,i+n);
        g.add_edge(v,i+n);
    }
    HLD hld(g);
    StaticTopTreeRerootingDP<decltype(hld),TreeDP> dp(hld);
    while(q--){
        int op;
        cin >> op;
        if(op==0){
            int u;
            mint x;
            cin >> u >> x;
            a[u]=x;
            dp.update(u);
        }else{
            int e;
            mint x,y;
            cin >> e >> x >> y;
            e+=n;
            b[e]=x,c[e]=y;
            dp.update(e);
        }
        int r;
        cin >> r;
        cout << dp.query_reroot(r).ans << "\n";
    }
}

#line 1 "verify/yosupo/data-structure/point_set_tree_path_composite_sum.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/point_set_tree_path_composite_sum"
#line 1 "template.hpp"
#include<bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

using ll = long long;
using db = long double;
using vi = vector<int>;
using vl = vector<ll>;
using vd = vector<db>;
using pii = pair<int,int>;
using pll = pair<ll,ll>;
using pdd = pair<db,db>;
const int INF=INT_MAX/2;
const int MOD=998244353;
const int MOD2=1000000007;
const ll LINF=LLONG_MAX/2;
const db DINF=numeric_limits<db>::infinity();
const db EPS=1e-9;
const db PI=acos(db(-1));

template<class T>
using ordered_set = tree<T,null_type,less<T>,rb_tree_tag,tree_order_statistics_node_update>;
template<class T>
using ordered_multiset = tree<T,null_type,less_equal<T>,rb_tree_tag,tree_order_statistics_node_update>;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
mt19937_64 rng64(chrono::steady_clock::now().time_since_epoch().count());
#line 2 "graph/graph-base.hpp"

/**
 * Author: Teetat T.
 * Date: 2024-06-15
 * Description: Graph Base
 */

template<class T>
struct Edge{
    int from,to,id;
    T cost;
    Edge(int _from,int _to,T _cost,int _id):from(_from),to(_to),cost(_cost),id(_id){}
    operator int()const{return to;}
};

template<class T=void,bool directed=false>
struct Graph{
    static constexpr bool is_directed=directed;
    static constexpr bool is_weighted=!is_same<T,void>::value;
    using cost_type = std::conditional_t<is_weighted,T,int>;
    using edge_type = Edge<cost_type>;
    int n,m;
    vector<edge_type> edges;
    vector<vector<edge_type>> g;
    vector<int> _deg,_indeg,_outdeg;
    Graph():n(0),m(0){}
    Graph(int _n):n(_n),m(0),g(_n){}
    vector<edge_type> &operator[](int u){return g[u];}
    void add_edge(int from,int to,cost_type cost=1,int id=-1){
        assert(0<=from&&from<n&&0<=to&&to<n);
        if(id==-1)id=m;
        edges.emplace_back(edge_type(from,to,cost,id));
        g[from].emplace_back(edge_type(from,to,cost,id));
        if(!is_directed)g[to].emplace_back(edge_type(to,from,cost,id));
        m++;
    }
    void calc_deg(){
        _deg.assign(n,0);
        for(auto &e:edges)_deg[e.from]++,_deg[e.to]++;
    }
    void calc_inout_deg(){
        _indeg.assign(n,0);
        _outdeg.assign(n,0);
        for(auto &e:edges)_outdeg[e.from]++,_indeg[e.to]++;
    }
    const vector<int> &deg_array(){
        if(_deg.empty())calc_deg();
        return _deg;
    }
    const vector<int> &indeg_array(){
        if(_indeg.empty())calc_inout_deg();
        return _indeg;
    }
    const vector<int> &outdeg_array(){
        if(_outdeg.empty())calc_inout_deg();
        return _outdeg;
    }
    int deg(int u){return deg_array()[u];}
    int indeg(int u){return indeg_array()[u];}
    int outdeg(int u){return outdeg_array()[u];}
    Graph reverse(){
        assert(is_directed);
        Graph res(n);
        for(auto &e:edges)res.add_edge(e.to,e.from,e.cost,e.id);
        return res;
    }
};

template<class T=void,bool directed=false>
Graph<T,directed> read_graph(int n,int m,int offset=1){
    using G = Graph<T,directed>;
    G g(n);
    for(int i=0;i<m;i++){
        int u,v;
        cin >> u >> v;
        u-=offset,v-=offset;
        if(g.is_weighted){
            typename G::cost_type w;
            cin >> w;
            g.add_edge(u,v,w);
        }else{
            g.add_edge(u,v);
        }
    }
    return g;
}
template<class T=void,bool directed=false>
Graph<T,directed> read_tree(int n,int offset=1){
    return read_graph<T,directed>(n,n-1,offset);
}

#line 2 "tree/hld.hpp"

/**
 * Author: Teetat T.
 * Date: 2024-06-15
 * Description: Heavy-Light Decomposition.
 */

template<class G>
struct HLD{
    int n;
    G &g;
    int root,timer;
    vector<int> par,sz,dep,hv,head,tin,tout,ord;
    HLD(G &_g,int _root=0)
        : n(_g.n),g(_g),root(_root),timer(-1),par(n,root),sz(n,1),
          dep(n),hv(n,-1),head(n),tin(n),tout(n),ord(n){
        par[0]=-1;
        dfs_sz(root);
        dfs_hld(root);
    }
    void dfs_sz(int u){
        for(int v:g[u])if(v!=par[u]){
            par[v]=u;
            dep[v]=dep[u]+1;
            dfs_sz(v);
            sz[u]+=sz[v];
            if(hv[u]==-1||sz[v]>sz[hv[u]])hv[u]=v;
        }
    }
    void dfs_hld(int u){
        tin[u]=++timer;
        ord[timer]=u;
        if(hv[u]!=-1){
            head[hv[u]]=head[u];
            dfs_hld(hv[u]);
        }
        for(int v:g[u])if(v!=par[u]&&v!=hv[u]){
            head[v]=v;
            dfs_hld(v);
        }
        tout[u]=timer;
    }
    vector<pair<int,int>> get_path(int u,int v,bool vertex,bool commutative=true){
        vector<pair<int,int>> up,down;
        while(head[u]!=head[v]){
            if(dep[head[u]]>dep[head[v]]){
                up.emplace_back(tin[head[u]],tin[u]);
                u=par[head[u]];
            }else{
                down.emplace_back(tin[head[v]],tin[v]);
                v=par[head[v]];
            }
        }
        if(dep[u]>dep[v])up.emplace_back(tin[v]+1,tin[u]),u=v;
        else if(u!=v)down.emplace_back(tin[u]+1,tin[v]),v=u;
        if(vertex)up.emplace_back(tin[u],tin[u]);
        reverse(down.begin(),down.end());
        if(!commutative)for(auto &[x,y]:up)swap(x,y);
        up.insert(up.end(),down.begin(),down.end());
        return up;
    }
    int lca(int u,int v){
        while(head[u]!=head[v]){
            if(dep[head[u]]>dep[head[v]])swap(u,v);
            v=par[head[v]];
        }
        return dep[u]<dep[v]?u:v;
    }
    int dist(int u,int v){
        return dep[u]+dep[v]-2*dep[lca(u,v)];
    }
    int jump(int u,int v,int k){
        int w=lca(u,v);
        int d=dep[u]+dep[v]-2*dep[w];
        if(k>d)return -1;
        if(k>dep[u]-dep[w]){
            k=d-k;
            swap(u,v);
        }
        while(k>=dep[u]-dep[head[u]]+1){
            k-=dep[u]-dep[head[u]]+1;
            u=par[head[u]];
        }
        return ord[tin[u]-k];
    }
    bool is_ancestor(int u,int v){
        return tin[u]<=tin[v]&&tout[v]<=tout[u];
    }
};

#line 2 "tree/static-top-tree.hpp"

/**
 * Author: Teetat T.
 * Date: 2024-11-14
 * Description: Static Top Tree.
 */

template<class HLD>
struct StaticTopTree{
    using P = pair<int,int>;
    enum Type{Compress,Rake,AddEdge,AddVertex,Vertex};
    int n,root;
    HLD &hld;
    vector<int> lch,rch,par;
    vector<Type> type;
    StaticTopTree(HLD &_hld):hld(_hld){build();}
    void build(){
        n=hld.n;
        lch=rch=par=vector<int>(n,-1);
        type.assign(n,Compress);
        root=compress(hld.root).second;
    }
    int add(int i,int l,int r,Type t){
        if(i==-1){
            i=n++;
            lch.emplace_back(l);
            rch.emplace_back(r);
            par.emplace_back(-1);
            type.emplace_back(t);
        }else{
            lch[i]=l,rch[i]=r,type[i]=t;
        }
        if(l!=-1)par[l]=i;
        if(r!=-1)par[r]=i;
        return i;
    }
    P compress(int i){
        vector<P> a{add_vertex(i)};
        auto work=[&](){
            auto [sj,j]=a.back();
            a.pop_back();
            auto [si,i]=a.back();
            a.back()={max(si,sj)+1,add(-1,i,j,Compress)};
        };
        while(hld.hv[i]!=-1){
            a.emplace_back(add_vertex(i=hld.hv[i]));
            while(true){
                if(a.size()>=3&&(a.end()[-3].first==a.end()[-2].first||a.end()[-3].first<=a.back().first)){
                    P tmp=a.back();
                    a.pop_back();
                    work();
                    a.emplace_back(tmp);
                }else if(a.size()>=2&&a.end()[-2].first<=a.back().first){
                    work();
                }else break;
            }
        }
        while(a.size()>=2)work();
        return a[0];
    }
    P rake(int i){
        priority_queue<P,vector<P>,greater<P>> pq;
        for(int j:hld.g[i])if(j!=hld.par[i]&&j!=hld.hv[i])pq.emplace(add_edge(j));
        while(pq.size()>=2){
            auto [si,i]=pq.top();pq.pop();
            auto [sj,j]=pq.top();pq.pop();
            pq.emplace(max(si,sj)+1,add(-1,i,j,Rake));
        }
        return pq.empty()?make_pair(0,-1):pq.top();
    }
    P add_edge(int i){
        auto [sj,j]=compress(i);
        return {sj+1,add(-1,j,-1,AddEdge)};
    }
    P add_vertex(int i){
        auto [sj,j]=rake(i);
        return {sj+1,add(i,j,-1,j==-1?Vertex:AddVertex)};
    }
};

#line 3 "tree/static-top-tree-rerooting-dp.hpp"

/**
 * Author: Teetat T.
 * Date: 2024-11-14
 * Description: Static Top Tree DP.
 */

/*
struct TreeDP{
    struct Path{
        static Path unit();
    };
    struct Point{
        static Point unit();
    };
    static Path compress(Path l,Path r);
    static Point rake(Point l,Point r);
    static Point add_edge(Path p);
    static Path add_vertex(Point p,int u);
    static Path vertex(int u);
};
*/

template<class HLD,class TreeDP>
struct StaticTopTreeRerootingDP{
    using Path = typename TreeDP::Path;
    using Point = typename TreeDP::Point;
    StaticTopTree<HLD> stt;
    vector<Path> path,rpath;
    vector<Point> point;
    StaticTopTreeRerootingDP(HLD &hld):stt(hld){
        int n=stt.n;
        path.resize(n);
        point.resize(n);
        rpath.resize(n);
        dfs(stt.root);
    }
    void _update(int u){
        if(stt.type[u]==stt.Vertex){
            path[u]=rpath[u]=TreeDP::vertex(u);
        }else if(stt.type[u]==stt.Compress){
            path[u]=TreeDP::compress(path[stt.lch[u]],path[stt.rch[u]]);
            rpath[u]=TreeDP::compress(rpath[stt.rch[u]],rpath[stt.lch[u]]);
        }else if(stt.type[u]==stt.Rake){
            point[u]=TreeDP::rake(point[stt.lch[u]],point[stt.rch[u]]);
        }else if(stt.type[u]==stt.AddEdge){
            point[u]=TreeDP::add_edge(path[stt.lch[u]]);
        }else{
            path[u]=rpath[u]=TreeDP::add_vertex(point[stt.lch[u]],u);
        }
    }
    void dfs(int u){
        if(u==-1)return;
        dfs(stt.lch[u]);
        dfs(stt.rch[u]);
        _update(u);
    }
    void update(int u){
        for(;u!=-1;u=stt.par[u])_update(u);
    }
    Path query_all(){
        return path[stt.root];
    }
    Path query_subtree(int u){
        Path res=path[u];
        while(true){
            int p=stt.par[u];
            if(p==-1||stt.type[p]!=stt.Compress)break;
            if(stt.lch[p]==u)res=TreeDP::compress(res,path[stt.rch[p]]);
        }
        return res;
    }
    Path query_reroot(int u){
        auto rec=[&](auto &&rec,int u)->Point{
            int p=stt.par[u];
            Path below=Path::unit(),above=Path::unit();
            while(p!=-1&&stt.type[p]==stt.Compress){
                int l=stt.lch[p],r=stt.rch[p];
                if(l==u)below=TreeDP::compress(below,path[r]);
                else above=TreeDP::compress(above,rpath[l]);
                u=p;
                p=stt.par[u];
            }
            if(p!=-1){
                u=p;
                p=stt.par[u];
                Point sum=Point::unit();
                while(stt.type[p]==stt.Rake){
                    int l=stt.lch[p],r=stt.rch[p];
                    sum=TreeDP::rake(sum,u==r?point[l]:point[r]);
                    u=p;
                    p=stt.par[u];
                }
                sum=TreeDP::rake(sum,rec(rec,p));
                above=TreeDP::compress(above,TreeDP::add_vertex(sum,p));
            }
            return TreeDP::rake(TreeDP::add_edge(below),TreeDP::add_edge(above));
        };
        Point res=rec(rec,u);
        if(stt.type[u]==stt.AddVertex){
            res=TreeDP::rake(res,point[stt.lch[u]]);
        }
        return TreeDP::add_vertex(res,u);
    }
};

#line 2 "modular-arithmetic/montgomery-modint.hpp"

/**
 * Author: Teetat T.
 * Date: 2024-03-17
 * Description: modular arithmetic operators using Montgomery space
 */

template<uint32_t mod,uint32_t root=0>
struct MontgomeryModInt{
    using mint = MontgomeryModInt;
    using i32 = int32_t;
    using u32 = uint32_t;
    using u64 = uint64_t;

    static constexpr u32 get_r(){
        u32 res=1;
        for(i32 i=0;i<5;i++)res*=2-mod*res;
        return res;
    }

    static const u32 r=get_r();
    static const u32 n2=-u64(mod)%mod;
    static_assert(mod<(1<<30));
    static_assert((mod&1)==1);
    static_assert(r*mod==1);

    u32 x;

    constexpr MontgomeryModInt():x(0){}
    constexpr MontgomeryModInt(const int64_t &v):x(reduce(u64(v%mod+mod)*n2)){}

    static constexpr u32 get_mod(){return mod;}
    static constexpr mint get_root(){return mint(root);}
    explicit constexpr operator int64_t()const{return val();}

    static constexpr u32 reduce(const u64 &v){
        return (v+u64(u32(v)*u32(-r))*mod)>>32;
    }

    constexpr u32 val()const{
        u32 res=reduce(x);
        return res>=mod?res-mod:res;
    }

    constexpr mint inv()const{
        int a=val(),b=mod,u=1,v=0,q=0;
        while(b>0){
            q=a/b;
            a-=q*b;
            u-=q*v;
            swap(a,b);
            swap(u,v);
        }
        return mint(u);
    }

    constexpr mint &operator+=(const mint &rhs){
        if(i32(x+=rhs.x-2*mod)<0)x+=2*mod;
        return *this;
    }
    constexpr mint &operator-=(const mint &rhs){
        if(i32(x-=rhs.x)<0)x+=2*mod;
        return *this;
    }
    constexpr mint &operator*=(const mint &rhs){
        x=reduce(u64(x)*rhs.x);
        return *this;
    }
    constexpr mint &operator/=(const mint &rhs){
        return *this*=rhs.inv();
    }

    constexpr mint &operator++(){return *this+=mint(1);}
    constexpr mint &operator--(){return *this-=mint(1);}
    constexpr mint operator++(int){
        mint res=*this;
        return *this+=mint(1),res;
    }
    constexpr mint operator--(int){
        mint res=*this;
        return *this-=mint(1),res;
    }

    constexpr mint operator-()const{return mint()-mint(*this);};
    constexpr mint operator+()const{return mint(*this);};

    friend constexpr mint operator+(const mint &lhs,const mint &rhs){return mint(lhs)+=rhs;}
    friend constexpr mint operator-(const mint &lhs,const mint &rhs){return mint(lhs)-=rhs;}
    friend constexpr mint operator*(const mint &lhs,const mint &rhs){return mint(lhs)*=rhs;}
    friend constexpr mint operator/(const mint &lhs,const mint &rhs){return mint(lhs)/=rhs;}
    friend constexpr bool operator==(const mint &lhs,const mint &rhs){
        return (lhs.x>=mod?lhs.x-mod:lhs.x)==(rhs.x>=mod?rhs.x-mod:rhs.x);
    }
    friend constexpr bool operator!=(const mint &lhs,const mint &rhs){
        return (lhs.x>=mod?lhs.x-mod:lhs.x)!=(rhs.x>=mod?rhs.x-mod:rhs.x);
    }
    friend constexpr bool operator<(const mint &lhs,const mint &rhs){
        return (lhs.x>=mod?lhs.x-mod:lhs.x)<(rhs.x>=mod?rhs.x-mod:rhs.x); // for std::map
    }

    friend istream &operator>>(istream &is,mint &o){
        int64_t v;
        is >> v;
        o=mint(v);
        return is;
    }
    friend ostream &operator<<(ostream &os,const mint &o){
        return os << o.val();
    }
};
using mint998 = MontgomeryModInt<998244353,3>;
using mint107 = MontgomeryModInt<1000000007>;

#line 7 "verify/yosupo/data-structure/point_set_tree_path_composite_sum.test.cpp"

using mint = mint998;

int n;
vector<int> id;
vector<mint> a,b,c;

struct TreeDP{
    struct Path{
        mint a,b,cnt,ans;
        static Path unit(){
            return {1,0,0,0};
        }
    };
    struct Point{
        mint cnt,ans;
        static Point unit(){
            return {0,0};
        }
    };
    static Path compress(const Path &l,const Path &r){
        Path res;
        res.a=l.a*r.a;
        res.b=l.a*r.b+l.b;
        res.cnt=l.cnt+r.cnt;
        res.ans=l.ans+l.a*r.ans+l.b*r.cnt;
        return res;
    }
    static Point rake(const Point &l,const Point &r){
        return {l.cnt+r.cnt,l.ans+r.ans};
    }
    static Point add_edge(const Path &p){
        return {p.cnt,p.ans};
    }
    static Path add_vertex(const Point &p,int u){
        Path res{b[u],c[u],p.cnt+(u<n),p.ans+a[u]};
        res.ans=res.a*res.ans+res.b*res.cnt;
        return res;
    }
    static Path vertex(int u){
        Path res{b[u],c[u],u<n,a[u]};
        res.ans=res.a*res.ans+res.b*res.cnt;
        return res;
    }
};

int main(){
    cin.tie(nullptr)->sync_with_stdio(false);
    int q;
    cin >> n >> q;
    a.resize(2*n);
    b.resize(2*n-1,mint(1));
    c.resize(2*n-1);
    for(int i=0;i<n;i++){
        cin >> a[i];
    }
    Graph<void,false> g(2*n-1);
    for(int i=0;i<n-1;i++){
        int u,v;
        cin >> u >> v >> b[i+n] >> c[i+n];
        g.add_edge(u,i+n);
        g.add_edge(v,i+n);
    }
    HLD hld(g);
    StaticTopTreeRerootingDP<decltype(hld),TreeDP> dp(hld);
    while(q--){
        int op;
        cin >> op;
        if(op==0){
            int u;
            mint x;
            cin >> u >> x;
            a[u]=x;
            dp.update(u);
        }else{
            int e;
            mint x,y;
            cin >> e >> x >> y;
            e+=n;
            b[e]=x,c[e]=y;
            dp.update(e);
        }
        int r;
        cin >> r;
        cout << dp.query_reroot(r).ans << "\n";
    }
}