結果

問題 No.439 チワワのなる木
ユーザー mugen_1337
提出日時 2020-07-21 23:20:23
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 328 ms / 5,000 ms
コード長 3,462 bytes
コンパイル時間 3,405 ms
コンパイル使用メモリ 210,704 KB
最終ジャッジ日時 2025-01-12 01:48:41
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 28
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include<bits/stdc++.h>
using namespace std;
#define ALL(x) x.begin(),x.end()
#define rep(i,n) for(int i=0;i<(n);i++)
#define debug(v) cout<<#v<<":";for(auto x:v){cout<<x<<' ';}cout<<endl;
#define mod 1000000007
using ll=long long;
const int INF=1000000000;
const ll LINF=1001002003004005006ll;
int dx[]={1,0,-1,0};
int dy[]={0,1,0,-1};
// ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
template<class T>bool chmax(T &a,const T &b){if(a<b){a=b;return true;}return false;}
template<class T>bool chmin(T &a,const T &b){if(b<a){a=b;return true;}return false;}

struct IOSetup{
    IOSetup(){
        cin.tie(0);
        ios::sync_with_stdio(0);
        cout<<fixed<<setprecision(12);
    }
} iosetup;

template<typename T1,typename T2>
ostream &operator<<(ostream &os,const pair<T1,T2>&p){
    os<<p.first<<" "<<p.second;
    return os;
}

template<typename T>
ostream &operator<<(ostream &os,const vector<T>&v){
    for(int i=0;i<(int)v.size();i++) os<<v[i]<<(i+1==v.size()?"":" ");
    return os;
}

template<typename T>
struct rerooting{
    using F=function<T(T,int)>;//集合，頂点番号
    using M=function<T(T,T)>;
 
    int V;
    vector<vector<int>> G;
    vector<vector<T>> dp;
    vector<T> ans;
    // dp_v = g(merge(f(dp_c1,c1),...,f(dp_ck,ck)),v)
    F f,g;  // TxN->T
    M merge;// TxT->T，子を集約する関数，モノイド
    T gen;
    //gen: mergeの元
    rerooting(int V,F f,M merge,T gen,F g=[](T a,int b){return a;})
        :V(V),f(f),merge(merge),gen(gen),g(g),G(V),dp(V),ans(V,gen){}
    
    //bidirectional
    void add_edge(int a,int b){
        G[a].push_back(b);
        G[b].push_back(a);
    }
    T dfs1(int pre,int now){
        T ret=gen;
        for(int i=0;i<G[now].size();i++)if(G[now][i]!=pre){
            dp[now][i]=dfs1(now,G[now][i]);
            ret=merge(ret,f(dp[now][i],G[now][i]));
        }
        return g(ret,now);
    }
    void dfs2(int pre,int now,T frompar){
        for(int i=0;i<G[now].size();i++){
            if(G[now][i]==pre){
                dp[now][i]=frompar;
                break;
            }
        }
        vector<T> lsum(G[now].size()+1),rsum(G[now].size()+1);//親も混ぜて累積
        lsum[0]=gen;rsum[G[now].size()]=gen;
        for(int i=0;i<G[now].size();i++) lsum[i+1]=merge(lsum[i],f(dp[now][i],G[now][i]));
        for(int i=G[now].size();i>0;i--) rsum[i-1]=merge(rsum[i],f(dp[now][i-1],G[now][i-1]));
        for(int i=0;i<G[now].size();i++)if(G[now][i]!=pre){
            T fromme=merge(lsum[i],rsum[i+1]);
            dfs2(now,G[now][i],g(fromme,now));
        }
    }
    void build(int root=0){
        for(int i=0;i<V;i++) dp[i].resize(G[i].size());
        dfs1(-1,root);
        dfs2(-1,root,gen);
        for(int i=0;i<V;i++){
            for(int j=0;j<G[i].size();j++) ans[i]=merge(ans[i],f(dp[i][j],G[i][j]));
            ans[i]=g(ans[i],i);
        }
    }
};

using T=tuple<ll,ll,ll>;

signed main(){
    int n;cin>>n;
    string s;cin>>s;
    auto f=[&](T t,int idx){return t;};
    auto merge=[&](T a,T b){return T(get<0>(a)+get<0>(b),get<1>(a)+get<1>(b),0);};
    auto g=[&](T a,int idx){
        auto [w,ww,cww]=a;
        if(s[idx]=='c') return T(w,ww,ww);
        else            return T(w+1,ww+w,cww);
    };

    rerooting<T> R(n,f,merge,T(0,0,0),g);
    rep(i,n-1){
        int u,v;cin>>u>>v;u--,v--;
        R.add_edge(u,v);
    }
    R.build();
    ll ans=0;
    rep(i,n)if(s[i]=='c')ans+=get<2>(R.ans[i]);
    cout<<ans<<endl;
    return 0;
}
0