test/rerooting.test.cpp

View this file on GitHub
Last update: 2024-04-25 14:16:22+09:00
Problem: https://judge.yosupo.jp/problem/tree_path_composite_sum

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/tree_path_composite_sum"


#include "../tree/rerooting.hpp"
#include "../math/mint998.hpp"


#include <bits/stdc++.h>
using namespace std;
using ll =long long;



using mint=modint<998244353>;

struct Data{
    ll b;
    ll c;
};



int main() {
    ll N;cin>>N;
    vector<ll> a(N);
    for(ll i=0;i<N;i++) {
        cin>>a[i];

    }

vector<pair<mint,mint>> A(N);
for(ll i=0;i<N;i++) {
    A[i]=make_pair(a[i],1);
}

auto f1=[](pair<mint,mint> x,pair<mint,mint> y) {
return make_pair(x.first+y.first,x.second+y.second);
};

auto f2=[](pair<mint,mint> x,Data y) {
    return make_pair(x.first*y.b+x.second*y.c,x.second);
};

ReRooting<Data,pair<mint,mint>> g(N,f1,f2,{0,0});
g.add_plus(A);

for(ll i=0;i<N-1;i++) {
    ll u,v,b,c;cin>>u>>v>>b>>c;
    g.add_edge(u,v,{b,c});
}

auto ans=g.solve();

for(auto x:ans) {
    cout<<x.first<<endl;
}
}

#line 1 "test/rerooting.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/tree_path_composite_sum"


#line 2 "tree/rerooting.hpp"

#include<bits/stdc++.h>
using namespace std;


/*
使い方
Data : 各頂点が持つもの
T：各辺が持つもの
F1：頂点同士をマージする時のマージの仕方
F2：子から親へと情報が移動するときの変化

*/

template< typename Data, typename T >
struct ReRooting {
 
  struct Node {
    int to, rev;
    Data data;
  };
 
  using F1 = function< T(T, T) >;
  using F2 = function< T(T, Data) >;
 
  vector< vector< Node > > g;
  vector< vector< T > > ldp, rdp;
  vector<T> plus;
  vector< int > lptr, rptr;
  const F1 f1;
  const F2 f2;
  const T ident;
 
  ReRooting(int n, const F1 &f1, const F2 &f2, const T &ident) :
      g(n), ldp(n), rdp(n), lptr(n), rptr(n), plus(n,ident),f1(f1), f2(f2), ident(ident) {}

 void add_plus(vector<T> plus_) {
  plus=plus_;
 }

  void add_edge(int u, int v, const Data &d) {
    g[u].emplace_back((Node) {v, (int) g[v].size(), d});
    g[v].emplace_back((Node) {u, (int) g[u].size() - 1, d});
  }
 
  void add_edge_bi(int u, int v, const Data &d, const Data &e) {
    g[u].emplace_back((Node) {v, (int) g[v].size(), d});
    g[v].emplace_back((Node) {u, (int) g[u].size() - 1, e});
  }
 
 
  T dfs(int idx, int par) {
 
    while(lptr[idx] != par && lptr[idx] < g[idx].size()) {
      auto &e = g[idx][lptr[idx]];
      ldp[idx][lptr[idx] + 1] = f1(ldp[idx][lptr[idx]], f2(dfs(e.to, e.rev), e.data));
      ++lptr[idx];
    }


    while(rptr[idx] != par && rptr[idx] >= 0) {
      auto &e = g[idx][rptr[idx]];
      rdp[idx][rptr[idx]] = f1(rdp[idx][rptr[idx] + 1], f2(dfs(e.to, e.rev), e.data));
      --rptr[idx];
    }

    if(par < 0) return f1(rdp[idx][0],plus[idx]);
    return f1(f1(ldp[idx][par], rdp[idx][par + 1]),plus[idx]);
  }
 
  vector< T > solve() {
    for(int i = 0; i < g.size(); i++) {
      ldp[i].assign(g[i].size() + 1, ident);
      rdp[i].assign(g[i].size() + 1, ident);
      lptr[i] = 0;
      rptr[i] = (int) g[i].size() - 1;
    }
    vector< T > ret;
    for(int i = 0; i < g.size(); i++) {
      ret.push_back(dfs(i, -1));
    }
    return ret;
  }
};
#line 3 "math/mint998.hpp"
using namespace std;
using ll =long long;

ll INF=2e18;


template<int mod> 
class modint {
    long long x;
public:
    modint(long long x=0) : x((x%mod+mod)%mod) {}
    modint operator-() const { 
      return modint(-x);
    }
    modint& operator+=(const modint& a) {
        if ((x += a.x) >= mod) x -= mod;
        return *this;
    }
    modint& operator-=(const modint& a) {
        if ((x += mod-a.x) >= mod) x -= mod;
        return *this;
    }
    modint& operator*=(const  modint& a) {
        (x *= a.x) %= mod;
        return *this;
    }
    modint operator+(const modint& a) const {
        modint res(*this);
        return res+=a;
    }
    modint operator-(const modint& a) const {
        modint res(*this);
        return res-=a;
    }
    modint operator*(const modint& a) const {
        modint res(*this);
        return res*=a;
    }
    modint pow(ll t) const {
        if (!t) return 1;
        modint a = pow(t>>1);
        a *= a;
        if (t&1) a *= *this;
        return a;
    }
    // for prime mod
    modint inv() const {
        return pow(mod-2);
    }
    modint& operator/=(const modint& a) {
        return (*this) *= a.inv();
    }
    modint operator/(const modint& a) const {
        modint res(*this);
        return res/=a;
    }

    bool operator==(const modint &a) const {
        modint res(*this);
        return res.x==a.x;
    }

    bool operator!=(const modint &a) const {
        modint res(*this);
        return res.x!=a.x;
    }


    friend ostream& operator<<(ostream& os, const modint& m){
        os << m.x;
        return os;
    }
};
#line 6 "test/rerooting.test.cpp"


#line 9 "test/rerooting.test.cpp"
using namespace std;
using ll =long long;



using mint=modint<998244353>;

struct Data{
    ll b;
    ll c;
};



int main() {
    ll N;cin>>N;
    vector<ll> a(N);
    for(ll i=0;i<N;i++) {
        cin>>a[i];

    }

vector<pair<mint,mint>> A(N);
for(ll i=0;i<N;i++) {
    A[i]=make_pair(a[i],1);
}

auto f1=[](pair<mint,mint> x,pair<mint,mint> y) {
return make_pair(x.first+y.first,x.second+y.second);
};

auto f2=[](pair<mint,mint> x,Data y) {
    return make_pair(x.first*y.b+x.second*y.c,x.second);
};

ReRooting<Data,pair<mint,mint>> g(N,f1,f2,{0,0});
g.add_plus(A);

for(ll i=0;i<N-1;i++) {
    ll u,v,b,c;cin>>u>>v>>b>>c;
    g.add_edge(u,v,{b,c});
}

auto ans=g.solve();

for(auto x:ans) {
    cout<<x.first<<endl;
}
}