TopCoder SRM 799 : Div1 Medium MapleTree

本番間に合わず…。
https://community.topcoder.com/stat?c=problem_statement&pm=15469&rd=18494

問題

木を成すN頂点の無向グラフが与えられる。
辺には長さが設定されている。
このうち、頂点をいくつか選んだとする。

選んだ頂点群を連結に保つ辺と頂点だけ残したグラフの直径を考える。
頂点の選びかた2^N-1通りのうち、直径の和はいくつか。

解法

まず直径の候補を列挙しよう。直径の中点を扱うので、辺の長さをO(N^2*logN)かけて頂点間の距離を求めつつ、2頂点の中点が、どこの点または辺になるかを求めておこう。
中点が辺に来るケースについては、辺の両端から、直径の両端までの距離のペアも覚えておく。
距離のペアが同じケースは、まとめて処理できる。

次に、各点または各辺に直径の中点が来るケースを考える。

• グラフの頂点が直径の中点になるケース
  • その点を木の根として、各子頂点の先のSubtreeにおける根からの距離から、Subtree内の頂点の選び方による最遠点とその組み合わせを求める。
  • すべての子頂点の先の最遠点までの距離がd以下で、かつ最遠点がちょうどdとなる子頂点が2個以上あるケースを数え上げる。
• 辺の途中が直径の中点になるケース
  • 辺の両端の先にSubtreeにおける根からの距離から、Subtree内の頂点の選び方による最遠点とその組み合わせを求める。
  • 最遠点までの距離が、事前に求めた直径の両端までの距離のペアに一致するものを数え上げる。

int N;
vector<pair<int,int>> E[2020];
const ll mo=1000000007;

ll dp[2020][2020];
vector<pair<int,int>> tar[2020][2020];

vector<pair<ll,int>> V;

void dfs(int cur,int pre,int id,ll d) {
	V.push_back({d,cur});
	dp[id][cur]=d;
	int x=lower_bound(ALL(V),make_pair(d/2,0))-V.begin();
	if(id<cur) {
		if(V[x].first!=d/2) {
			int a=V[x-1].second;
			int b=V[x].second;
			if(a<b) tar[a][b].push_back({id,cur});
			else tar[b][a].push_back({cur,id});
		}
		else {
			tar[V[x].second][V[x].second].push_back({cur,id});
		}
	}
	FORR(e,E[cur]) if(e.first!=pre) dfs(e.first,cur,id,d+e.second);
	V.pop_back();
}

void dfs2(int cur,int pre,vector<ll>& A,ll d) {
	A.push_back(d);
	FORR(e,E[cur]) if(e.first!=pre) dfs2(e.first,cur,A,d+e.second);
	
}
class MapleTree {
	public:
	int sum(vector <int> p, vector <int> length) {
		N=p.size()+1;
		int i,j,x,y;
		
		FOR(i,N+1) E[i].clear();
		FOR(x,N) FOR(y,N) tar[x][y].clear();
		FOR(i,p.size()) {
			E[p[i]].push_back({i+1,length[i]*2});
			E[i+1].push_back({p[i],length[i]*2});
		}
		FOR(i,N) dfs(i,i,i,0);
		ll ret=0;
		FOR(i,N) {
			set<ll> cand;
			FORR(a,tar[i][i]) {
				cand.insert(dp[a.first][a.second]/2);
			}
			vector<vector<ll>> A;
			FORR(e,E[i]) {
				vector<ll> B;
				dfs2(e.first,i,B,e.second);
				sort(ALL(B));
				A.push_back(B);
			}
			FORR(c,cand) {
				ll dp[3]={1,0,0};
				FORR(b,A) {
					ll v=1;
					ll le=1,eq=0;
					FORR(a,b) {
						if(a<c) (le+=v)%=mo;
						if(a==c) (eq+=v)%=mo;
						v=v*2%mo;
					}
					ll dp2[3]={dp[0],dp[1],dp[2]};
					dp[0]=dp2[0]*le%mo;
					dp[1]=(dp2[1]*le+dp2[0]*eq)%mo;
					dp[2]=(dp2[1]*eq+dp2[2]*(le+eq))%mo;
				}
				(ret+=dp[2]*c*2*2)%=mo;
			}
		}
		FOR(x,N) FOR(y,N) if(x!=y&&tar[x][y].size()) {
			vector<ll> A,B;
			dfs2(x,y,A,0);
			dfs2(y,x,B,0);
			sort(ALL(A));
			sort(ALL(B));
			map<ll,ll> MA,MB;
			ll v=1;
			FORR(a,A) (MA[a]+=v)%=mo, v=v*2%mo;
			v=1;
			FORR(a,B) (MB[a]+=v)%=mo, v=v*2%mo;

			set<pair<ll,ll>> cand;
			FORR(e,tar[x][y]) cand.insert({dp[e.first][x],dp[y][e.second]});
			FORR(e,cand) {
				ll s=e.first,t=e.second;
				ll len=s+t+dp[x][y];
				(ret+=MA[s]*MB[t]%mo*len)%=mo;
			}
			
		}
		
		return ret*(mo+1)/2%mo;
	}
}

まとめ

本番も方針はあってたんだけど、実装が間に合わなかった。
そのままだと微妙にTLEしてたので、もともと本番だと解き切れてなかったことになるな。