This documentation is automatically generated by online-judge-tools/verification-helper
View the Project on GitHub beet-aizu/library
#define PROBLEM "https://yukicoder.me/problems/3961" #include<bits/stdc++.h> using namespace std; #define call_from_test #include "../../tree/centroid.cpp" #include "../../tools/fixpoint.cpp" #include "../../vector/compress.cpp" #undef call_from_test #ifdef SANITIZE #define IGNORE #endif signed main(){ cin.tie(0); ios::sync_with_stdio(0); using ll = long long; int n,k; cin>>n>>k; using P = pair<int, int>; vector< map<int, int> > col(n); Centroid G(n); for(int i=1;i<n;i++){ int x,y,z; cin>>x>>y>>z; x--;y--;z--; G.add_edge(x,y); col[x][y]=z; col[y][x]=z; } queue<int> que; que.emplace(G.build(0)[0]); ll ans=0; vector<int> cnt(k,0),uku(k,0); vector<P> ps(n),vp(n); while(!que.empty()){ int r=que.front();que.pop(); auto calc= [&](P p,int q)->P{ assert(~p.first); if(p.first==q||p.second==q) return p; if(p.second==-1) return minmax({p.first,q}); return P(-1,-1); }; vp.clear(); for(int t:G[r]){ if(!G.alive(t)) continue; MFP([&](auto dfs,int v,int p,P st)->void{ ps[v]=st; if(st.first<0) return; for(int u:G[v]){ if(u==p) continue; if(!G.alive(u)) continue; dfs(u,v,calc(st,col[u][v])); } if(st.second!=-1) vp.emplace_back(st); })(t,r,P(col[r][t],-1)); } vp=compress(vp); auto idx=[&](P p){return lower_bound(vp.begin(),vp.end(),p)-vp.begin();}; vector<int> num(vp.size()); int all=0; for(int t:G[r]){ if(!G.alive(t)) continue; // count MFP([&](auto dfs,int v,int p)->void{ P st=ps[v]; if(st.first<0) return; for(int u:G[v]){ if(u==p) continue; if(!G.alive(u)) continue; dfs(u,v); } if(st.second==-1){ ans+=all-cnt[st.first]; ans+=uku[st.first]; }else{ ans+=1; ans+=cnt[st.first]; ans+=cnt[st.second]; ans+=num[idx(st)]; } })(t,r); // paint MFP([&](auto dfs,int v,int p)->void{ P st=ps[v]; if(st.first<0) return; for(int u:G[v]){ if(u==p) continue; if(!G.alive(u)) continue; dfs(u,v); } if(st.second==-1){ cnt[st.first]++; all++; }else{ num[idx(st)]++; uku[st.first]++; uku[st.second]++; } })(t,r); } for(int t:G[r]){ if(!G.alive(t)) continue; // clear MFP([&](auto dfs,int v,int p)->void{ P st=ps[v]; if(st.first<0) return; for(int u:G[v]){ if(u==p) continue; if(!G.alive(u)) continue; dfs(u,v); } if(st.second==-1){ cnt[st.first]--; }else{ uku[st.first]--; uku[st.second]--; } })(t,r); } G.disable(r); for(int u:G[r]) if(G.alive(u)) que.emplace(G.build(u)[0]); } cout<<ans<<endl; return 0; }
#line 1 "test/yukicoder/3961.test.cpp" #define PROBLEM "https://yukicoder.me/problems/3961" #include<bits/stdc++.h> using namespace std; #define call_from_test #line 1 "tree/centroid.cpp" #line 3 "tree/centroid.cpp" using namespace std; #endif //BEGIN CUT HERE struct Centroid{ vector<int> sz,dead; vector< vector<int> > G; Centroid(int n):sz(n,1),dead(n,0),G(n){} void add_edge(int u,int v){ G[u].emplace_back(v); G[v].emplace_back(u); } int dfs(int v,int p){ sz[v]=1; for(int u:G[v]) if(u!=p and !dead[u]) sz[v]+=dfs(u,v); return sz[v]; } void find(int v,int p,int tmp,vector<int> &cs) { int ok=1; for (int u:G[v]){ if(u==p or dead[u]) continue; find(u,v,tmp,cs); ok&=(sz[u]<=tmp/2); } ok&=(tmp-sz[v]<=tmp/2); if(ok) cs.emplace_back(v); } vector<int> build(int r) { int tmp=dfs(r,-1); vector<int> cs; find(r,-1,tmp,cs); return cs; } const vector<int>& operator[](int k)const{return G[k];} void disable(int v){dead[v]=1;} void enable(int v){dead[v]=0;} int alive(int v){return !dead[v];} }; //END CUT HERE #ifndef call_from_test signed main(){ return 0; } #endif #line 1 "tools/fixpoint.cpp" #line 3 "tools/fixpoint.cpp" using namespace std; #endif //BEGIN CUT HERE template<typename F> struct FixPoint : F{ FixPoint(F&& f):F(forward<F>(f)){} template<typename... Args> decltype(auto) operator()(Args&&... args) const{ return F::operator()(*this,forward<Args>(args)...); } }; template<typename F> inline decltype(auto) MFP(F&& f){ return FixPoint<F>{forward<F>(f)}; } //END CUT HERE #ifndef call_from_test //INSERT ABOVE HERE signed main(){ return 0; } #endif #line 1 "vector/compress.cpp" #line 3 "vector/compress.cpp" using namespace std; #endif //BEGIN CUT HERE template<typename V> V compress(V vs){ sort(vs.begin(),vs.end()); vs.erase(unique(vs.begin(),vs.end()),vs.end()); return vs; } template<typename T> map<T, int> dict(const vector<T> &vs){ map<T, int> res; for(int i=0;i<(int)vs.size();i++) res[vs[i]]=i; return res; } map<char, int> dict(const string &s){ return dict(vector<char>(s.begin(),s.end())); } template<typename T> vector<T> compressed(vector<T> vs){ auto dc=dict(compress(vs)); for(auto &v:vs) v=dc[v]; return vs; } //END CUT HERE #ifndef call_from_test //INSERT ABOVE HERE signed main(){ return 0; } #endif #line 10 "test/yukicoder/3961.test.cpp" #undef call_from_test #ifdef SANITIZE #define IGNORE #endif signed main(){ cin.tie(0); ios::sync_with_stdio(0); using ll = long long; int n,k; cin>>n>>k; using P = pair<int, int>; vector< map<int, int> > col(n); Centroid G(n); for(int i=1;i<n;i++){ int x,y,z; cin>>x>>y>>z; x--;y--;z--; G.add_edge(x,y); col[x][y]=z; col[y][x]=z; } queue<int> que; que.emplace(G.build(0)[0]); ll ans=0; vector<int> cnt(k,0),uku(k,0); vector<P> ps(n),vp(n); while(!que.empty()){ int r=que.front();que.pop(); auto calc= [&](P p,int q)->P{ assert(~p.first); if(p.first==q||p.second==q) return p; if(p.second==-1) return minmax({p.first,q}); return P(-1,-1); }; vp.clear(); for(int t:G[r]){ if(!G.alive(t)) continue; MFP([&](auto dfs,int v,int p,P st)->void{ ps[v]=st; if(st.first<0) return; for(int u:G[v]){ if(u==p) continue; if(!G.alive(u)) continue; dfs(u,v,calc(st,col[u][v])); } if(st.second!=-1) vp.emplace_back(st); })(t,r,P(col[r][t],-1)); } vp=compress(vp); auto idx=[&](P p){return lower_bound(vp.begin(),vp.end(),p)-vp.begin();}; vector<int> num(vp.size()); int all=0; for(int t:G[r]){ if(!G.alive(t)) continue; // count MFP([&](auto dfs,int v,int p)->void{ P st=ps[v]; if(st.first<0) return; for(int u:G[v]){ if(u==p) continue; if(!G.alive(u)) continue; dfs(u,v); } if(st.second==-1){ ans+=all-cnt[st.first]; ans+=uku[st.first]; }else{ ans+=1; ans+=cnt[st.first]; ans+=cnt[st.second]; ans+=num[idx(st)]; } })(t,r); // paint MFP([&](auto dfs,int v,int p)->void{ P st=ps[v]; if(st.first<0) return; for(int u:G[v]){ if(u==p) continue; if(!G.alive(u)) continue; dfs(u,v); } if(st.second==-1){ cnt[st.first]++; all++; }else{ num[idx(st)]++; uku[st.first]++; uku[st.second]++; } })(t,r); } for(int t:G[r]){ if(!G.alive(t)) continue; // clear MFP([&](auto dfs,int v,int p)->void{ P st=ps[v]; if(st.first<0) return; for(int u:G[v]){ if(u==p) continue; if(!G.alive(u)) continue; dfs(u,v); } if(st.second==-1){ cnt[st.first]--; }else{ uku[st.first]--; uku[st.second]--; } })(t,r); } G.disable(r); for(int u:G[r]) if(G.alive(u)) que.emplace(G.build(u)[0]); } cout<<ans<<endl; return 0; }