【今日の精進】ABC163F - path pass i

毎日格上の問題を倒すやつの93日目です。

問題リンク

atcoder.jp

問題概要

頂点に色がついた木があるよ。

各色について、その色を一回以上通るパスの個数を求めてね。

解法

頂点0を根とした根つき木とします。

トポロジカルソートをして、オイラーツアー順を取得しておきます。これはどちらもDFSで対応できます。

各色、ソート順が大きい方から、「その頂点を通って、以前に処理した頂点たちは通らないパスの総数」を求めたいです。

これは、先ほど作ったオイラーツアーから遅延セグメント木を作成すれば、部分木の削除と部分木の頂点数の取得を高速にできます。

よって、これらからいい感じに計算すればよいです。

提出コード

#include<bits/stdc++.h>
#include<atcoder/lazysegtree>
using namespace std;
using namespace atcoder;
using ll = long long;
#define rep(i,n) for(int i=0; i<n; i++)
#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")

void dfs(int pos,int bef,int &cnt,
        vector<int> &c,vector<vector<int>> &tree,vector<vector<int>> &col,vector<pair<int,int>> &tour){
    tour[pos].first=cnt;
    cnt++;
    col[c[pos]-1].push_back(pos);
    for(int i:tree[pos]){
        if(i!=bef){
            dfs(i,pos,cnt,c,tree,col,tour);
        }
    }
    tour[pos].second=cnt;
    cnt++;
}

pair<int,int> op(pair<int,int> x,pair<int,int> y){
    return pair<int,int>(x.first+y.first,x.second+y.second);
}

pair<int,int> e(){
    return pair<int,int>(0,0);
}

pair<int,int> mapp(int f,pair<int,int> x){
    if(f==-1)return x;
    return pair<int,int>(x.second*f,x.second);
}

int comp(int f,int g){
    if(f==-1)return g;
    return f;
}

int id(){
    return -1;
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int N;
    cin>>N;
    vector<int> c(N);
    rep(i,N){
        cin>>c[i];
    }
    vector<vector<int>> tree(N);
    rep(i,N-1){
        int a,b;
        cin>>a>>b;
        tree[a-1].push_back(b-1);
        tree[b-1].push_back(a-1);
    }
    vector<vector<int>> col(N);
    vector<pair<int,int>> tour(N);
    int cnt=0;
    dfs(0,-1,cnt,c,tree,col,tour);
    lazy_segtree<pair<int,int>,op,e,int,mapp,comp,id> seg(vector<pair<int,int>>(N*2,pair<int,int>(1,1)));
    rep(i,N){
        seg.apply(0,N*2,1);
        reverse(col[i].begin(),col[i].end());
        ll ans=0;
        for(int j:col[i]){
            ll add=0,squ=0;
            for(int k:tree[j]){
                if(tour[j].first<tour[k].first){
                    ll size=seg.prod(tour[k].first,tour[k].second+1).first/2;
                    add+=size;
                    squ+=size*size;
                }
            }
            ll size=(seg.all_prod().first-seg.prod(tour[j].first,tour[j].second+1).first)/2;
            add+=size;
            squ+=size*size;
            ans+=(add*add-squ)/2+seg.all_prod().first/2;
            seg.apply(tour[j].first,tour[j].second+1,0);
        }
        cout<<ans<<endl;
    }
}

atcoder.jp

感想

オイラーツアーだけでなくいろんな典型が使えるいい問題でした