Submission #501274

#	Time	Username	Problem	Language	Result	Execution time	Memory
501274		rainboy	Wells (CEOI21_wells)	C11	100 / 100	7297 ms	406948 KiB

wells

#include <stdio.h>

#define N	1500000
#define INF	0x3f3f3f3f3f3f3f3fLL
#define MD	1000000007
#define H_	22	/* H_ = ceil(log2(N * 2)) */
#define N_	(1 << H_)

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

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

int ll[N], jj[N], ta[N], tb[N], kk[N + 1], qu[N * 2];

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

	kk[d + 1]++;
	ll[i] = 0, jj[i] = -1;
	for (o = eo[i]; o--; ) {
		int j = ej[i][o];

		if (j != p) {
			dfs1(i, j, d + 1);
			if (ll[i] < ll[j])
				ll[i] = ll[j], jj[i] = j;
		}
	}
	ll[i]++;
}

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

	qu[ta[i] = time++] = i, tb[i] = kk[d]++;
	if (jj[i] != -1)
		dfs2(i, jj[i], d + 1);
	else
		time++;
	for (o = eo[i]; o--; ) {
		int j = ej[i][o];

		if (j != p && j != jj[i])
			dfs2(i, j, d + 1);
	}
}

int cnt[N * 2]; long long xx[N + 1];

void get_range(int a, int *l, int *r) {
	*l = tb[qu[a]], *r = *l + cnt[a];
}

void dfs3(int p, int i) {
	int o, d, l1, r1, l2, r2;
	long long c;

	for (o = eo[i]; o--; ) {
		int j = ej[i][o];

		if (j != p)
			dfs3(i, j);
	}
	if (k <= ll[i]) {
		get_range(ta[i] + k - 1, &l1, &r1);
		xx[l1]++, xx[r1]--;
		if (p != -1)
			xx[tb[p]] -= r1 - l1, xx[tb[p] + 1] += r1 - l1;
	}
	cnt[ta[i]]++;
	for (o = eo[i]; o--; ) {
		int j = ej[i][o];

		if (j != p && j != jj[i]) {
			for (d = 1; d <= ll[j] && d < k; d++)
				if (k - d <= ll[i]) {
					get_range(ta[i] + k - d - 1, &l1, &r1), get_range(ta[j] + d - 1, &l2, &r2);
					c = (long long) (r1 - l1) * (r2 - l2);
					xx[l1] += r2 - l2, xx[r1] -= r2 - l2;
					xx[l2] += r1 - l1, xx[r2] -= r1 - l1;
					xx[tb[i]] -= c, xx[tb[i] + 1] += c;
					if (p != -1)
						xx[tb[p]] -= c, xx[tb[p] + 1] += c;
				}
			for (d = 1; d <= ll[j]; d++)
				cnt[ta[i] + d] += cnt[ta[j] + d - 1];
		}
	}
}

long long ww[N];

void dfs4(int p, int i) {
	int o;

	ww[i] = xx[tb[i]];
	for (o = eo[i]; o--; ) {
		int j = ej[i][o];

		if (j != p) {
			dfs4(i, j);
			ww[i] += ww[j];
		}
	}
}

int h_, n_;
long long stu[N_ * 2], lzu[N_]; int stv[N_ * 2], lzv[N_];

void put(int i, long long u, long long v) {
	if (stu[i] != INF && u != INF)
		stu[i] += u, stv[i] = (long long) stv[i] * v % MD;
	else
		stu[i] = INF, stv[i] = 0;
	if (i < n_) {
		if (lzu[i] != INF && u != INF)
			lzu[i] += u, lzv[i] = (long long) lzv[i] * v % MD;
		else
			lzu[i] = INF, lzv[i] = 0;
	}
}

void pus(int i) {
	if (lzu[i] != 0 || lzv[i] != 1)
		put(i << 1 | 0, lzu[i], lzv[i]), put(i << 1 | 1, lzu[i], lzv[i]), lzu[i] = 0, lzv[i] = 1;
}

void pul(int i) {
	if (lzu[i] == 0 && lzv[i] == 1) {
		int l = i << 1, r = l | 1;

		if (stu[l] < stu[r])
			stu[i] = stu[l], stv[i] = stv[l];
		else if (stu[l] > stu[r])
			stu[i] = stu[r], stv[i] = stv[r];
		else
			stu[i] = stu[l], stv[i] = (stv[l] + stv[r]) % MD;
	}
}

void push(int i) {
	int h;

	for (h = h_; h > 0; h--)
		pus(i >> h);
}

void pull(int i) {
	while (i > 1)
		pul(i >>= 1);
}

void upd(int i, long long u, int v) {
	push(i += n_);
	stu[i] = u, stv[i] = v;
	pull(i);
}

void update(int l, int r, long long u, int v) {
	if (l <= r) {
		int l_ = l += n_, r_ = r += n_;

		push(l_), push(r_);
		for ( ; l <= r; l >>= 1, r >>= 1) {
			if ((l & 1) == 1)
				put(l++, u, v);
			if ((r & 1) == 0)
				put(r--, u, v);
		}
		pull(l_), pull(r_);
	}
}

void query(int l, int r, long long *u_, int *v_) {
	long long u = INF, v = 0;

	if (l <= r) {
		push(l += n_), push(r += n_);
		for ( ; l <= r; l >>= 1, r >>= 1) {
			if ((l & 1) == 1) {
				if (u > stu[l])
					u = stu[l], v = 0;
				if (u == stu[l])
					v = (v + stv[l]) % MD;
				l++;
			}
			if ((r & 1) == 0) {
				if (u > stu[r])
					u = stu[r], v = 0;
				if (u == stu[r])
					v = (v + stv[r]) % MD;
				r--;
			}
		}
	}
	*u_ = u, *v_ = v;
}

void process(int i, int j, int d) {
	long long u1, ui, uj, u;
	int v1, vi, vj, v;

	u1 = INF, v1 = 0;
	query(ta[i] + d, ta[i] + d, &ui, &vi);
	query(ta[j], ta[j] + min(ll[j], min(k - d - 1, d - 2)), &uj, &vj);
	u = ui + uj;
	v = (long long) vi * vj % MD;
	if (u1 > u)
		u1 = u, v1 = 0;
	if (u1 == u)
		v1 = (v1 + v) % MD;
	if (d - 1 <= ll[j]) {
		query(ta[j] + d - 1, ta[j] + d - 1, &uj, &vj);
		query(ta[i] + 1, ta[i] + min(k - d, d), &ui, &vi);
		u = ui + uj;
		v = (long long) vi * vj % MD;
		if (u1 > u)
			u1 = u, v1 = 0;
		if (u1 == u)
			v1 = (v1 + v) % MD;
	}
	upd(ta[i] + d, u1, v1);
}

void dfs5(int p, int i) {
	int o, j, d;
	long long u;
	int v;

	upd(ta[i] + 0, ww[i], 1);
	if (jj[i] == -1)
		upd(ta[i] + 1, 0, 1);
	else {
		for (o = eo[i]; o--; ) {
			j = ej[i][o];
			if (j != p)
				dfs5(i, j);
		}
		j = jj[i];
		query(ta[j], ta[j] + ll[j], &u, &v);
		update(ta[i] + 0, ta[i] + 0, u, v);
		if (ll[i] >= k)
			upd(ta[i] + k, INF, 0);
		for (o = eo[i]; o--; ) {
			j = ej[i][o];
			if (j != p && j != jj[i]) {
				int lower, upper;

				query(ta[j], ta[j] + ll[j], &u, &v);
				update(ta[i] + 0, ta[i] + 0, u, v);
				lower = ll[j] + 2, upper = min(ll[i], k - 1 - ll[j]);
				if (lower > upper)
					for (d = min(ll[i], k); d > 0; d--)
						process(i, j, d);
				else {
					for (d = min(ll[i], k); d > upper; d--)
						process(i, j, d);
					query(ta[j], ta[j] + ll[j], &u, &v);
					update(ta[i] + lower, ta[i] + upper, u, v);
					for (d = lower - 1; d > 0; d--)
						process(i, j, d);
				}
			}
		}
	}
}

int main() {
	static int ii_[N - 1], jj_[N - 1];
	int h, i, j, d, v;
	long long w, u;

	scanf("%d%d", &n, &k);
	for (h = 0; h < n - 1; h++) {
		scanf("%d%d", &i, &j), i--, j--;
		ii_[h] = i, jj_[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 = ii_[h], j = jj_[h];
		ej[i][eo[i]++] = j, ej[j][eo[j]++] = i;
	}
	dfs1(-1, 0, 0);
	for (d = 1; d <= n; d++)
		kk[d] += kk[d - 1];
	dfs2(-1, 0, 0);
	dfs3(-1, 0);
	for (i = 1; i <= n; i++)
		xx[i] += xx[i - 1];
	dfs4(-1, 0);
	w = 0;
	for (i = 0; i < n; i++)
		w += ww[i];
	w /= k;
	h_ = 0;
	while (1 << h_ < n * 2)
		h_++;
	n_ = 1 << h_;
	for (i = 0; i < n_; i++)
		lzv[i] = 1;
	dfs5(-1, 0);
	query(ta[0], ta[0] + ll[0], &u, &v);
	if (u == w)
		printf("YES\n%d\n", v);
	else
		printf("NO\n0\n");
	return 0;
}

Compilation message (stderr)

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

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