結果

問題 No.3194 Do Optimize Your Solution
ユーザー Rubikun
提出日時 2025-06-28 04:03:28
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
TLE  
実行時間 -
コード長 9,819 bytes
コンパイル時間 3,341 ms
コンパイル使用メモリ 229,300 KB
実行使用メモリ 215,748 KB
最終ジャッジ日時 2025-06-28 04:04:17
合計ジャッジ時間 14,111 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 12 TLE * 5
権限があれば一括ダウンロードができます

ソースコード

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;


vector<int> G2[MAX];

int vs[MAX*2];
int depth[MAX*2];
int id[MAX];

void dfs(int v,int p,int d,int &k){
    id[v]=k;
    vs[k]=v;
    depth[k++]=d;
    for(int to:G2[v]){
        if(to==p) continue;
        dfs(to,v,d+1,k);
        vs[k]=v;
        depth[k++]=d;
    }
    //cout<<k<<" ";
}

void init(){
    int k=0;
    dfs(0,-1,0,k);
}

struct SparseTable{
    using T=int;
    //using F=function<T(T,T)>;
    
    int n,logn;
    vector<vector<T>> dat;
    vector<int> loglen;
    vector<int> depthsv;
    
    SparseTable(int n_){
        n=1;
        logn=1;
        while(n<n_){
            n*=2;
            logn++;
        }
        loglen.resize(n+3);
        depthsv.resize(n+3,INF);
        n=n_;
        int j=0;
        for(int i=1;i<n+3;i++){
            loglen[i]=j;
            if(i+1>(1<<(j+1))) j++;
        }
        
        dat.resize(logn);
        
        for(int i=0;i<logn;i++){
            dat[i].assign(n+1,n+2);
        }
    }
    
    void set(int j,pair<int,int> x){
        auto [d,id]=x;
        dat[0][j]=id;
        depthsv[id]=d;
    }
    
    void built(){
        for(int i=1;i<logn;i++){
            for(int j=0;j<n+1;j++){
                T vla=dat[i-1][j],vr;
                if(j+(1<<(i-1))>=n+1) vr=n+2;
                else vr=dat[i-1][j+(1<<(i-1))];
                if(depthsv[vla]<depthsv[vr]) dat[i][j]=vla;
                else dat[i][j]=vr;
            }
        }
    }
    
    int query(int l,int r){
        if(l>=r){
            return -1;
        }else{
            T vla=dat[loglen[r-l]][l],vr=dat[loglen[r-l]][r-(1<<loglen[r-l])];
            if(depthsv[vla]<depthsv[vr]){
                return vla;
            }else{
                return vr;
            }
        }
    }
};

SparseTable spa(1);

int lca(int a,int b){
    return spa.query(min(id[a],id[b]),max(id[a],id[b])+1);
}

int dd(int a,int b){
    return depth[id[a]]+depth[id[b]]-2*depth[id[lca(a,b)]];
}

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]];
    }
};

struct AuxiliaryTree : LowestCommonAncestor{
    using super = LowestCommonAncestor;
    
    vector<int> idx,cn,omo;
    vector<ll> dp;
    vvii T;
    AuxiliaryTree(int n):super(n),idx(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){
        int d=dd(u,v);
        T[u].emplace_back(mp(v,d));
        T[v].emplace_back(mp(u,d));
    }
    
    using 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();
        vi st;
        st.reserve(2*k-1);
        st.pb(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.back();st.pop_back();
                while(!st.empty() and dep[w]<dep[st.back()]){
                    add_aux_edge(st.back(),l);
                    l=st.back();st.pop_back();
                }
                if(st.empty() or st.back()!=w){
                    st.pb(w);
                    vs.emplace_back(w);
                }
                add_aux_edge(w,l);
            }
            st.pb(vs[i+1]);
        }
        
        while(st.size()>1){
            int c=st.back();st.pop_back();
            add_aux_edge(st.back(),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]+=(ll)cn[to]*dd;
            cn[u]+=cn[to];
        }
    }
    
    void solve(int u,int p){
        dp[u]=0;
        for(auto [to,dd]:T[u]){
            dp[u]+=dp[to];
            dp[u]+=(ll)cn[to]*dd;
        }
        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]-=(ll)cn[to]*dd;
            cn[u]-=cn[to];
            cn[to]+=cn[u];
            
            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;
            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);
    }
}

vi use;

vvii Y;
int tim;
array<int,3> pos[MAX][20];
int poscur[MAX];

void calc(vii S){
    if(si(S)<=1) return;
    
    for(int j=0;j<si(S);j++){
        auto [a,b]=S[j];
        pos[AT.idx[b]][poscur[AT.idx[b]]++]={tim,a,b};
        //pos[AT.idx[b]].pb({tim,a,b});
    }
    tim++;
    
    return;
}

int W[MAX];

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;
    S.reserve(subtree_size[s]);
    S.pb(mp(0,s));
    
    vi T={1};
    for(int a:G[s]){
        if(centroid[a]){
            continue;
        }
        atume(a,s,1,S);
        T.pb(si(S));
    }
    
    for(int j=0;j<si(S);j++){
        S[j].fi+=W[S[j].se];
    }
    calc(S);
    for(int j=0;j<si(S);j++){
        S[j].fi-=W[S[j].se];
        W[S[j].se]=0;
    }
    
    int q=1;
    
    for(int a:G[s]){
        if(centroid[a]){
            continue;
        }
        
        for(int j=T[q-1];j<T[q];j++){
            W[S[j].se]-=S[j].fi;
        }
        solve_subproblem(a,s);
        q++;
    }
    
    centroid[s]=0;
}

void solve(){
    Y.resize(tim);
    vi CNY(tim);
    for(int i=0;i<N;i++){
        for(int j=0;j<poscur[i];j++){
            auto [t,a,b]=pos[i][j];
            CNY[t]++;
        }
    }
    for(int i=0;i<tim;i++){
        Y[i].reserve(CNY[i]);
    }
    for(int i=0;i<N;i++){
        for(int j=0;j<poscur[i];j++){
            auto [t,a,b]=pos[i][j];
            Y[t].pb(mp(a,b));
        }
    }
    for(int i=0;i<si(Y);i++){
        auto S=Y[i];
        
        vi use(si(S));
        for(int j=0;j<si(S);j++) use[j]=S[j].se;
        
        AT.query(use);
        
        for(auto [a,b]:S){
            AT.omo[b]=a;
            AT.cn[b]=1;
        }
        AT.make(use[0],-1);
        AT.solve(use[0],-1);
        for(int x:use){
            AT.dp[x]=0;
            AT.omo[x]=0;
            AT.cn[x]=0;
        }
        AT.clear(use);
    }
}

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--;
        G2[a].pb(b);
        G2[b].pb(a);
        AT.add_edge(a,b);
    }
    
    AT.build();
    
    init();
    
    spa=SparseTable(2*N-1);
    for(int i=0;i<2*N-1;i++){
        spa.set(i,mp(depth[i],vs[i]));
    }
    spa.built();
    
    solve_subproblem(0,-1);
    solve();
    
    ans*=2;
    cout<<ans<<endl;
}
0