日立製作所 社会システム事業部 プログラミングコンテスト2020 : F - Preserve Diameter

初手を間違ってしまった…。
https://atcoder.jp/contests/hitachi2020/tasks/hitachi2020_f

問題

木を成す無向グラフGが与えられる。
ここで辺のない2点間にいくつか辺を加えたグラフをHとする。
Hが以下を満たすような辺の追加の組み合わせは何通りか。

• GとHの直径は等しい。
• さらに辺を1本でも加えると、Hの直径が短くなってしまう。

解法

Gの直径をDとする。
このときHが満たすべき条件を考える。

• Hの直径を成す両端点がx,yとすると、x-y間以外に距離Dを成す頂点対はない。
  • 仮にxから距離Dであるzがあったとする。この時、距離D-2の点とzをつなぐとx-zの距離は短くなるがx-yの距離は変わらないので、条件に反する。
• よって、x,y以外の点はxからの距離が1～(D-1)である。
• xからvの距離をf(v)とする。|f(u)-f(v)|≦1であるような点はすでに辺がすべて張られている。

こうすると、辺が張られていない頂点対はxからの距離が2以上離れているので、辺を追加すると直径が縮まってしまう。

ここまで踏まえて、以下Dが偶数の時を考える。
x,yがどこになるかは候補が複数あるが、いずれにせよx-yの中点はグラフの中心Cと一致するはずである。
そこで、
f(v) := 符号付きのCからの距離
とし、辺に0,+1,-1の長さを振って、f(x)=D/2、f(y)=-D/2となるx,yが1個ずつあるようなケースを数え上げよう。
こうすると、頂点Cを根とする根付き木を考えたとき、
dp(v,w,a,b) := 頂点v以下のsubtreeにおいて、f(v)=wで、subtree以下にf(v')=D/2となる点がa個、f(v')=-D/2となる点がb個あるような組み合わせ
を求めていけばいいことになる。
ただ、f(x)=D/2やf(y)=-D/2となるx,yをsubtreeに抱えるには、頂点Cからそこに至るすべての辺が+1か-1に統一されなければならない。
そう考えるとほとんどのwにおいてa,bは0しか取りえないので、考えるべき状態を減らすことができる。
最終的にdp(C,0,1,1)を答えればよい。

Dが奇数の場合、中心が2点(u,v)になるが、同様にdp(u,0,0,1)*dp(v,0,1,0)を答えればよい。

int N;
vector<int> E[202020];
const ll mo=998244353;
ll dp[202020][3][3];

pair<int,int> farthest(int cur,int pre,int d,vector<int>& D) {
	D[cur]=d;
	pair<int,int> r={d,cur};
	FORR(e,E[cur]) if(e!=pre) r=max(r, farthest(e,cur,d+1,D));
	return r;
}

pair<int,vector<int>> diameter() { // diameter,center
	vector<int> D[2];
	D[0].resize(N);
	D[1].resize(N);
	auto v1=farthest(0,0,0,D[0]);
	auto v2=farthest(v1.second,v1.second,0,D[0]);
	farthest(v2.second,v2.second,0,D[1]);
	pair<int,vector<int>> R;
	R.first = v2.first;
	//重心を取る場合
	for(int i=N-1;i>=0;i--) if(D[0][i]+D[1][i]==R.first && abs(D[0][i]-D[1][i])<=1) R.second.push_back(i);

	return R;
}

void dfs(int cur,int pre,int lef) {
	if(lef==0) {
		dp[cur][1][1]=1;
	}
	else {
		dp[cur][0][0]=1;
	}
	
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur,lef-1);
		ll to[3][3]={};
		int a1,a2,b1,b2;
		FOR(a1,3) FOR(a2,3) FOR(b1,3) FOR(b2,3) {
			// +1
			to[min(a1+b1,2)][a2]+=dp[cur][a1][a2]*dp[e][b1][b2]%mo;
			// 0
			to[a1][a2]+=dp[cur][a1][a2]*dp[e][b1][b2]%mo;
			// -1
			to[a1][min(a2+b2,2)]+=dp[cur][a1][a2]*dp[e][b1][b2]%mo;
		}
		
		FOR(a1,3) FOR(a2,3) dp[cur][a1][a2]=to[a1][a2]%mo;
	}
	
}


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);
	}
	
	auto R=diameter();
	if(R.second.size()==1) {
		x=R.second[0];
		dfs(x,-1,R.first/2);
		cout<<dp[x][1][1]*((mo+1)/2)%mo<<endl;
	}
	else {
		x=R.second[0];
		y=R.second[1];
		dfs(x,y,R.first/2);
		dfs(y,x,R.first/2);
		ll X=(dp[x][0][1]+dp[x][1][1]+dp[x][2][1])%mo;
		ll Y=(dp[y][0][1]+dp[y][1][1]+dp[y][2][1])%mo;
		cout<<X*Y%mo<<endl;
	}
}

まとめ

本番何を勘違いしたか、Hにおいて距離Dの頂点対は1組でなくてもよいと思ってしまいどうにもならなかった。