考察をうまくやると一気に簡単になる。
http://codeforces.com/contest/1060/problem/E
問題
木を成すグラフがある。
このグラフに対し、「最初のグラフにおいて直接つながる辺はないが共通する隣接点を持つ頂点間に辺を追加する」という処理を行う。
その後のグラフにおいて、全頂点対の最短距離の総和を求めよ。
解法
追加の辺により、もともと距離2の点同士の間が距離1でつながるようになるので、基本的には(距離の総和)/2と答えたい。
ただ、もともと距離が奇数である頂点間は追加の辺により距離はceil(元の距離/2)となる。
そこで、元のグラフを2色で塗り分けよう。
そうすると距離が奇数となる頂点対の組み合わせは両色の頂点数の積となる。
こうすると、*1/2とシンプルな解になる。
全頂点対の距離の総和は全方位DPとかしなくてもO(N)で容易に算出できる。
これは定番なのでさっと掛けるようにしておいた方がよいね。
int N; vector<int> E[201010]; ll C[202020]; int DC[2]; ll tot; int dfs(int cur,int pre,int d) { DC[d]++; C[cur]++; FORR(e,E[cur]) if(e!=pre) { C[cur]+=dfs(e,cur,d^1); } tot+=C[cur]*(N-C[cur]); return C[cur]; } void solve() { int i,j,k,l,r,x,y; string s; cin>>N; FOR(i,N-1) { cin>>x>>y; E[x-1].push_back(y-1); E[y-1].push_back(x-1); } dfs(0,-1,0); tot+=1LL*DC[0]*DC[1]; cout<<tot/2<<endl; }
まとめ
最初愚直に距離の総和を求めようとしちゃった。
*1:元のグラフの距離の総和)+(上記頂点対の組み合わせ