Official

F - Diameter set Editorial by en_translator


Let \(T\) be the tree and \(d(u,v)\) denote the distance of \(u\) and \(v\) on \(T\).

Choose an arbitrary vertex, and let \(X\) be one of the furthest vertices from the vertex. Additionally, choose one of the furthest vertices from \(X\), denoting it by \(Y\). Then, it is known that \(d(X,Y)\) is equal to the diameter. Let \(D\) be the length. Also, we consider the path between \(X\) and \(Y\), denoting each vertex on the path (including the endpoints) by \(V_0(=X)\) , \(V_1\) , \(\ldots\) , \(V_D(=Y)\).

  • If \(D\) is odd

Define \(A\) and \(B\) by \(A=V_{\frac{D-1}{2}}\) and \(B=V_{\frac{D+1}{2}}\). Cut the edge between \(A\) and \(B\) and separate \(T\) into two subtrees \(T'\) and \(T"\).
Then, for any vertex \(v'\) in \(T'\), \(d(v',Y)=d(v',A)+d(A,B)+d(B,Y)=d(v',A)+1+\frac{D-1}{2}\leq D\), so \(d(A,v')\leq \frac{D-1}{2}\). Thus, for any two vertices \(u'\) and \(v'\) belonging to \(T'\), \(d(u',v')\leq d(A,u')+d(A,v')\leq D-1\). Similarly, for any two vertices \(u"\) and \(v"\) belonging to \(T"\), \(d(u",v")\leq d(B,u")+d(B,v")\leq D-1\).
Therefore, to meet the conditions, we have to choose exactly one vertex from each of \(T'\) and \(T"\). The distance between two vertices \(v\) and \(v"\), taken from \(T'\) and \(T"\) respectively, is represented as \(d(v',v")=d(A,v')+d(A,B)+d(B,v")\), and since \(d(A,v')\leq \frac{D-1}{2}\) and \(d(B,v")\leq \frac{D-1}{2}\), \(d(v',v")=D\) if and only if \(d(A,v')= \frac{D-1}{2}\) and \(d(B,v")= \frac{D-1}{2}\).
Hence,, the answer is \(M_1M_2\), where \(M_1\) denotes the number of vertices of \(T'\) such that the distance from \(A\) is \(\frac{D-1}{2}\), and \(M_2\) denotes the number of vertices of \(T"\) such that the distance from \(A\) is \(\frac{D-1}{2}\)

  • If \(D\) is even

It can be solved with the similar discussion to the odd case. This time, we take \(C=V_{\frac{D}{2}}\) and define \(A_1\), \(\ldots\) , \(A_K\) to be the vertices directly connected to \(C\). Next, we remove \(C\) and all the edges connected to \(C\), so that \(T\) is split to \(K\) subtrees \(T_1\), \(\dots\), and \(T_K\), where \(T_i\) is defined to be the subtree which \(A_i\) belongs.
Moreover, since \(X\) and \(Y\) belongs to different subtrees, for any \(i\), at least one of \(X\) and \(Y\) does not belong to \(T_i\), so for any vertex \(v\) belonging to \(T_i\), one of \(d(X,v)\) and \(d(Y,v)\) can be expressed as \(\frac{D}{2}+d(C,A_i)+d(A_i,v)\), which is at most \(D\), so it is concluded that \(d(A_i,v)\leq \frac{D}{2}-1\).
Therefore, any two vertices belonging to each \(T_i\) is at most \(D-2\), so when the conditions are satisfied, at most one vertex in each subtree is painted red.
Two vertices \(u\) and \(v\) belonging to different subtrees \(T_i\) and \(T_j\) can be written as \(d(A_i,u)+d(A_i,A_j)+d(A_j,v)=d(A_i,u)+d(A_j,v)+2\), and since \(d(A_i,u)\leq \frac{D}{2}-1\) and \(d(A_j,v)\leq \frac{D}{2}-1\), \(d(u,v)=D\) if and only if \(d(A_i,u)= \frac{D}{2}-1\) and \(d(A_j,v)= \frac{D}{2}-1\).
Therefore, denoting by \(M_i\) the number of vertices in \(T_i\) such that the distance from \(A_i\) is \(\frac{D}{2}-1\), there are \((M_1+1)\cdots (M_K+1)\) ways of choosing at most \(1\) appropriate vertex from each tree. Subtracting one way of choosing no vertex and \((M_1+M_2\cdots +M_K)\) ways of choosing only one vertex, the answer is \((M_1+1)\cdots (M_K+1)-(M_1+M_2\cdots +M_K)-1\).

Finding the diameter, finding \(A\) and \(B\) or \(C\), counting the number of vertices satisfying the conditions (i.e. computing \(M_i\)), and the final part of computing the answer, can be processed in a total of \(O(N)\) time each, so the problem has been solved.

#include <bits/stdc++.h>
using namespace std;

#define N 200010
#define MOD (ll)998244353
#define ll long long
#define pb push_back
#define rep(i, n) for(ll i = 0; i < n; ++i)

vector<int>e[N];

int d[N][5];
int mx[5];
int mv[5];
int idx, blc, num, cnt;

void dfs(int k, int p) {
	if (p == -1)d[k][idx] = 0;
	else d[k][idx] = d[p][idx] + 1;
	if (mx[idx] < d[k][idx]) {
		mx[idx] = d[k][idx];
		mv[idx] = k;
	}
	if (d[k][idx] == num)cnt++;
	int sz = e[k].size();
	rep(i, sz) {
		if ((e[k][i] != blc) && (d[e[k][i]][idx] < 0))dfs(e[k][i], k);
	}
	return;
}

int main(void) {
	int n;
	int u, v, sz;
	vector<int>a;
	ll x, ans;

	rep(i, 5) {
		rep(j, N)d[j][i] = -1;
		mx[i] = -1;
	}
	blc = -1;

	cin >> n;
	rep(i, n - 1) {
		cin >> u >> v;
		e[u - 1].pb(v - 1);
		e[v - 1].pb(u - 1);
	}

	idx = 0;
	dfs(0, -1);
	idx = 1;
	dfs(mv[0], -1);
	idx = 2;
	dfs(mv[1], -1);

	if (mx[1] % 2 == 1) {
		rep(i, n) {
			if ((d[i][1] == (mx[1] / 2)) && (d[i][2] == (mx[1] / 2) + 1))u = i;
			if ((d[i][2] == (mx[1] / 2)) && (d[i][1] == (mx[1] / 2) + 1))v = i;
		}

		idx = 3;
		num = mx[1] / 2;

		blc = v;
		cnt = 0;
		dfs(u, -1);
		a.pb(cnt);

		blc = u;
		cnt = 0;
		dfs(v, -1);
		a.pb(cnt);
	}
	else {
		rep(i, n) {
			if ((d[i][1] == (mx[1] / 2)) && (d[i][2] == (mx[1] / 2)))u = i;
		}

		idx = 3;
		num = (mx[1] / 2) - 1;
		blc = u;

		sz = e[u].size();
		rep(i, sz) {
			cnt = 0;
			dfs(e[u][i], -1);
			a.pb(cnt);
		}
	}

	sz = a.size();
	ans = 1;
	rep(i, sz) {
		x = a[i] + 1;
		ans = (ans*x) % MOD;
	}
	rep(i, sz) {
		x = a[i];
		ans = (ans + MOD - x) % MOD;
	}
	ans = (ans + MOD - 1) % MOD;

	cout << ans << endl;
	return 0;
}

posted:
last update: