Submission #504054

#	Time	Username	Problem	Language	Result	Execution time	Memory
504054		rainboy	Wells (CEOI21_wells)	C11	100 / 100	2793 ms	212856 KiB

wells

/* https://hsin.hr/ceoi/competition/ceoi2021_day2_editorial.pdf */
#include <stdio.h>

#define N	1500000
#define MD	1000000007
#define INF	0x3f3f3f3f

int min(int a, int b) { return a < b ? a : b; }
int max(int a, int b) { return a > b ? a : b; }

int pp2[N + 1];

void init() {
	int i;

	pp2[0] = 1;
	for (i = 1; i <= N; i++)
		pp2[i] = pp2[i - 1] * 2 % MD;
}

int mem[(N - 1) * 2], *ptr = mem;
int *ej[N], eo[N], pp[N], n, k; char type[N];

int i_, l;
int ii[N]; char path[N];

void dfs1(int p, int i, int d) {
	int o;

	pp[i] = p;
	d++;
	if (l < d)
		l = d, i_ = i;
	for (o = eo[i]; o--; ) {
		int j = ej[i][o];

		if (j != p)
			dfs1(i, j, d);
	}
}

int gg[N], dd[N], hh[N];

void dfs2(int p, int i, int g, int d) {
	int o;

	gg[i] = g, dd[i] = d, hh[i] = 0;
	for (o = eo[i]; o--; ) {
		int j = ej[i][o];

		if (j != p && !path[j]) {
			dfs2(i, j, g, d + 1);
			hh[i] = max(hh[i], hh[j] + 1);
		}
	}
}

int dfs3(int p, int i, int d) {
	int o, x;

	if (d == 0 || type[i] == -1)
		return 1;
	x = 1;
	for (o = eo[i]; o--; ) {
		int j = ej[i][o];

		if (j != p)
			x = (long long) x * dfs3(i, j, d - 1) % MD;
	}
	return (x + 1) % MD;
}

int cc[N], dp[N], dq[N], dr[N];

int dfs4(int p, int i) {
	int o, p1, p2;

	dp[i] = -INF, dq[i] = 0, dr[i] = -1, p1 = p2 = -INF;
	for (o = eo[i]; o--; ) {
		int j = ej[i][o];

		if (j != p && !path[j] && type[j] == 1) {
			if (!dfs4(i, j))
				return 0;
			dp[i] = max(dp[i], dp[j]);
			if (p1 < dp[j])
				p2 = p1, p1 = dp[j];
			else if (p2 < dp[j])
				p2 = dp[j];
			if (cc[i] == 0 && dq[i] + dq[j] + 1 >= k - 1)
				return 0;
			dq[i] = max(dq[i], dq[j] + 1);
			if (dr[j] != -1)
				dr[i] = dr[j] + 1;
		}
	}
	if (cc[i] == 1)
		dp[i] = hh[i], dq[i] = -1, dr[i] = 0;
	else if (p2 != -INF && dr[i] * 2 <= k - 1 && dr[i] * 2 + p1 + p2 >= k - 1)
		return 0;
	return 1;
}

char bad[N]; int xx[N];

void solve() {
	static int aa[N + 1], yy[N], zz[N];
	int g, r, i, d;

	for (g = 0; g < l; g++)
		if (type[ii[g]] == 1) {
			int lower, upper;

			lower = g + 1, upper = min(g * 2, k - 1);
			if (lower <= upper)
				aa[lower]++, aa[upper + 1]--;
			lower = max(g * 2 - l + 1, l - k), upper = g - 1;
			if (lower <= upper) {
				lower %= k, upper %= k;
				aa[lower]++, aa[upper + 1]--;
				if (lower > upper)
					aa[0]++;
			}
		}
	for (r = 1; r < k; r++)
		aa[r] += aa[r - 1];
	for (r = 0; r < k; r++)
		if (aa[r])
			bad[r] = 1;
	for (i = 0; i < n; i++)
		if (!path[i] && type[i] == 1) {
			int r1, r2;

			g = gg[i], d = dd[i];
			r1 = (g + d) % k, r2 = (g - d % k + k) % k;
			if (r1 != r2) {
				if (g >= r1 && g <= (l - 1 - r1) / k * k + r1)
					bad[r1] = 1;
				if (g >= r2 && g <= (l - 1 - r2) / k * k + r2)
					bad[r2] = 1;
			}
		}
	for (r = 0; r < k; r++)
		yy[r] = zz[r] = 1;
	for (i = 0; i < n; i++)
		if (!path[i] && type[i] == 0 && (path[pp[i]] || type[pp[i]] == 1)) {
			g = gg[i], d = dd[i];
			if (g + d + hh[i] >= k - 1) {
				if (g + 1 < k)
					yy[g + 1] = (long long) yy[g + 1] * dfs3(pp[i], i, k - 1 - (g + d)) % MD;
			} else {
				if (g > l - k)
					zz[g - 1 - l + k] = (long long) zz[g - 1 - l + k] * dfs3(pp[i], i, k - 1 - (l - 1 - g + d)) % MD;
			}
		}
	for (r = 1; r < k; r++)
		yy[r] = (long long) yy[r] * yy[r - 1] % MD;
	for (r = k - 2; r >= 0; r--)
		zz[r] = (long long) zz[r] * zz[r + 1] % MD;
	for (r = 0; r < k; r++)
		xx[r] = (long long) yy[r] * zz[(r - l % k + k) % k] % MD;
}

void check(int r) {
	int g, i, d;

	for (i = 0; i < n; i++) {
		g = gg[i], d = dd[i];
		if (path[i])
			cc[i] = g % k == r ? 1 : 0;
		else if (type[i] == 1)
			cc[i] = g >= r ? ((g + d) % k == r ? 1 : 0) : ((g - d % k + k) % k == r ? 1 : 0);
	}
	for (g = 0; g < l; g++)
		if (type[ii[g]] == 1 && !dfs4(-1, ii[g])) {
			bad[r] = 1;
			return;
		}
}

int main() {
	static int uu[N - 1], vv[N - 1];
	int g, h, i, j, u, v, d, r, c, deep, yes, ans;

	init();
	scanf("%d%d", &n, &k);
	if (k == 2) {
		printf("YES\n");
		printf("2\n");
		return 0;
	}
	for (h = 0; h < n - 1; h++) {
		scanf("%d%d", &i, &j), i--, j--;
		uu[h] = i, vv[h] = j;
		eo[i]++, eo[j]++;
	}
	for (i = 0; i < n; i++) {
		ej[i] = ptr, ptr += eo[i];
		eo[i] = 0;
	}
	for (h = 0; h < n - 1; h++) {
		i = uu[h], j = vv[h];
		ej[i][eo[i]++] = j, ej[j][eo[j]++] = i;
	}
	i_ = -1, l = -1, dfs1(-1, 0, 0), u = i_;
	i_ = -1, l = -1, dfs1(-1, u, 0), v = i_;
	if (l < k) {
		printf("YES\n");
		printf("%d\n", pp2[n]);
		return 0;
	}
	for (i = v, d = l - 1; i != u; i = pp[i], d--)
		path[ii[d] = i] = 1;
	path[ii[0] = u] = 1;
	for (g = 0; g < l; g++)
		dfs2(-1, ii[g], g, 0);
	c = 0;
	for (i = 0; i < n; i++) {
		g = gg[i];
		if (max(g, l - 1 - g) + dd[i] + hh[i] < k - 1) {
			type[i] = -1;
			if (!path[i])
				c++;
		} else if (min(g, l - 1 - g) + dd[i] + hh[i] >= k - 1)
			type[i] = 1;
		else
			type[i] = 0;
	}
	solve();
	deep = 0;
	for (i = 0; i < n; i++)
		if (dd[i] * 2 >= k - 1) {
			deep = 1;
			break;
		}
	if (deep)
		for (r = 0; r < k; r++)
			if (!bad[r])
				check(r);
	yes = 0, ans = 0;
	for (r = 0; r < k; r++)
		if (!bad[r])
			yes = 1, ans = (ans + xx[r]) % MD;
	ans = (long long) ans * pp2[c] % MD;
	printf(yes ? "YES\n" : "NO\n");
	printf("%d\n", ans);
	return 0;
}

Compilation message (stderr)

wells.c: In function 'main':
wells.c:186:2: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
  186 |  scanf("%d%d", &n, &k);
      |  ^~~~~~~~~~~~~~~~~~~~~
wells.c:193:3: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
  193 |   scanf("%d%d", &i, &j), i--, j--;
      |   ^~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 30 / 30
• Subtask 2: 20 / 20
• Subtask 3: 20 / 20
• Subtask 4: 30 / 30