結果

問題 No.3194 Do Optimize Your Solution
ユーザー Rubikun
提出日時 2025-06-27 22:27:23
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
TLE  
実行時間 -
コード長 6,908 bytes
コンパイル時間 3,160 ms
コンパイル使用メモリ 231,644 KB
実行使用メモリ 86,636 KB
最終ジャッジ日時 2025-06-27 22:27:35
合計ジャッジ時間 11,566 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 1 TLE * 1 -- * 15
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long ll;
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; }
#define vi vector<int>
#define vl vector<ll>
#define vii vector<pair<int,int>>
#define vll vector<pair<ll,ll>>
#define vvi vector<vector<int>>
#define vvl vector<vector<ll>>
#define vvii vector<vector<pair<int,int>>>
#define vvll vector<vector<pair<ll,ll>>>
#define vst vector<string>
#define pii pair<int,int>
#define pll pair<ll,ll>
#define pb push_back
#define all(x) (x).begin(),(x).end()
#define mkunique(x) sort(all(x));(x).erase(unique(all(x)),(x).end())
#define fi first
#define se second
#define mp make_pair
#define si(x) int(x.size())
const int mod=998244353,MAX=200005,INF=15<<26;

ll ans=0;

//auxiliary-tree

// https://beet-aizu.github.io/library/tree/auxiliarytree.cpp.html
struct LowestCommonAncestor{
    int h;
    vector< vector<int> > G,par;
    vector<int> dep;
    LowestCommonAncestor(int n):G(n),dep(n){
        h=1;
        while((1<<h)<=n) h++;
        par.assign(h,vector<int>(n,-1));
    }
    
    void add_edge(int u,int v){
        G[u].emplace_back(v);
        G[v].emplace_back(u);
    }
    
    void dfs(int v,int p,int d){
        par[0][v]=p;
        dep[v]=d;
        for(int u:G[v])
            if(u!=p) dfs(u,v,d+1);
    }
    
    void build(int r=0){
        int n=G.size();
        dfs(r,-1,0);
        for(int k=0;k+1<h;k++)
            for(int v=0;v<n;v++)
                if(~par[k][v])
                    par[k+1][v]=par[k][par[k][v]];
    }
    
    int lca(int u,int v){
        if(dep[u]>dep[v]) swap(u,v);
        for(int k=0;k<h;k++)
            if((dep[v]-dep[u])>>k&1)
                v=par[k][v];
        
        if(u==v) return u;
        
        for(int k=h-1;k>=0;k--)
            if(par[k][u]!=par[k][v])
                u=par[k][u],v=par[k][v];
        
        return par[0][u];
    }
    
    int distance(int u,int v){
        return dep[u]+dep[v]-dep[lca(u,v)]*2;
    }
};

struct AuxiliaryTree : LowestCommonAncestor{
    using super = LowestCommonAncestor;
    
    vector<int> idx,iru;
    vl cn,omo,dp;
    vvll T;
    AuxiliaryTree(int n):super(n),idx(n),iru(n),cn(n),omo(n),dp(n),T(n){}
    
    void dfs(int v,int p,int &pos){
        idx[v]=pos++;
        for(int u:G[v])
            if(u!=p) dfs(u,v,pos);
    }
    
    void build(int r=0){
        super::build(r);
        int pos=0;
        dfs(r,-1,pos);
    }
    
    void add_aux_edge(int u,int v){
        T[u].emplace_back(mp(v,distance(u,v)));
        T[v].emplace_back(mp(u,distance(u,v)));
    }
    
    using super::lca, super::dep;
    void query(vector<int> &vs){
        assert(!vs.empty());
        sort(vs.begin(),vs.end(),
             [&](int a,int b){return idx[a]<idx[b];});
        vs.erase(unique(vs.begin(),vs.end()),vs.end());
        
        int k=vs.size();
        stack<int> st;
        st.emplace(vs[0]);
        for(int i=0;i+1<k;i++){
            int w=lca(vs[i],vs[i+1]);
            if(w!=vs[i]){
                int l=st.top();st.pop();
                while(!st.empty() and dep[w]<dep[st.top()]){
                    add_aux_edge(st.top(),l);
                    l=st.top();st.pop();
                }
                if(st.empty() or st.top()!=w){
                    st.emplace(w);
                    vs.emplace_back(w);
                }
                add_aux_edge(w,l);
            }
            st.emplace(vs[i+1]);
        }
        
        while(st.size()>1){
            int c=st.top();st.pop();
            add_aux_edge(st.top(),c);
        }
    }
    
    void make(int u,int p){
        for(auto [to,dd]:T[u]){
            if(to==p) continue;
            make(to,u);
            
            dp[u]+=dp[to];
            dp[u]+=cn[to]*dd;
            cn[u]+=cn[to];
        }
    }
    
    void solve(int u,int p){
        if(iru[u]) cn[u]=1;
        else cn[u]=0;
        dp[u]=0;
        for(auto [to,dd]:T[u]){
            dp[u]+=dp[to];
            dp[u]+=cn[to]*dd;
            cn[u]+=cn[to];
        }
        ans+=dp[u]*omo[u];
        //cout<<u<<" "<<dp[u]<<" "<<cn[u]<<" "<<omo[u]<<endl;
        
        for(auto [to,dd]:T[u]){
            if(to==p) continue;
            
            ll a=dp[u],b=cn[u],c=dp[to],d=cn[to];
            
            dp[u]-=dp[to];
            dp[u]-=cn[to]*dd;
            cn[u]-=cn[to];
            
            solve(to,u);
            
            dp[u]=a;
            cn[u]=b;
            dp[to]=c;
            cn[to]=d;
        }
    }
    
    void clear(const vector<int> &ws){
        for(int w:ws){
            cn[w]=0;
            omo[w]=0;
            dp[w]=0;
            iru[w]=0;
            T[w].clear();
        }
    }
};

AuxiliaryTree AT(1);

//重心分解(使える版)

int N,C=-1;
vi G[MAX];

bool centroid[MAX];
int subtree_size[MAX];
int centpar[MAX];

int compute_subtree_size(int u,int p){
    int c=1;
    for(int a:G[u]){
        if(a==p||centroid[a]) continue;
        c+=compute_subtree_size(a,u);
    }
    return subtree_size[u]=c;
}

pair<int,int> search_centroid(int u,int p,int t){
    pair<int,int> res={INF,-1};
    int s=1,m=0;
    for(int a:G[u]){
        if(a==p||centroid[a]) continue;
        
        res=min(res,search_centroid(a,u,t));
        
        m=max(m,subtree_size[a]);
        s+=subtree_size[a];
    }
    m=max(m,t-s);
    res=min(res,{m,u});
    return res;
}

void atume(int u,int p,int d,vii &T){
    T.pb(mp(d,u));
    for(int to:G[u]){
        if(to==p||centroid[to]) continue;
        atume(to,u,d+1,T);
    }
}

void calc(vii S,int kei){
    if(si(S)<=1) return;
    vi use;
    for(auto [a,b]:S) use.pb(b);
    AT.query(use);
    for(auto [a,b]:S){
        AT.iru[b]=1;
        AT.omo[b]=a*kei;
        AT.cn[b]=1;
    }
    AT.make(use[0],-1);
    AT.solve(use[0],-1);
    AT.clear(use);
}

void solve_subproblem(int u,int p){
    compute_subtree_size(u,-1);
    int s=search_centroid(u,-1,subtree_size[u]).second;
    centroid[s]=1;
    if(C==-1) C=s;
    centpar[s]=p;
    
    vii S={mp(0,s)};
    
    for(int a:G[s]){
        if(centroid[a]){
            continue;
        }
        vii T;
        atume(a,s,1,T);
        calc(T,-1);
        for(auto x:T) S.pb(x);
        solve_subproblem(a,s);
    }
    
    calc(S,1);
    
    centroid[s]=0;
}

int main(){
    
    std::ifstream in("text.txt");
    std::cin.rdbuf(in.rdbuf());
    cin.tie(0);
    ios::sync_with_stdio(false);
    
    cin>>N;
    for(int i=0;i<N-1;i++){
        int a,b;cin>>a>>b;
        a--;b--;
        G[a].pb(b);
        G[b].pb(a);
    }
    
    AT=AuxiliaryTree(N);
    
    for(int i=0;i<N-1;i++){
        int a,b;cin>>a>>b;
        a--;b--;
        AT.add_edge(a,b);
    }
    
    AT.build();
    
    solve_subproblem(0,-1);
    ans*=2;
    cout<<ans<<endl;
}
0