Submission #882848

#	Time	Username	Problem	Language	Result	Execution time	Memory
882848		rainboy	Balanced Tree (info1cup18_balancedtree)	C11	100 / 100	880 ms	81124 KiB

balancedtree

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define N	500000
#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 n, d;

int *ej[N], eo[N], cc[N];

void append(int i, int j) {
	int o = eo[i]++;

	if (o >= 2 && (o & o - 1) == 0)
		ej[i] = (int *) realloc(ej[i], o * 2 * sizeof *ej[i]);
	ej[i][o] = j;
}

int dp[N][2][2][2], dq[N][2][2][2];

void dfs1(int p, int i) {
	int k, o, o_, c, c1, c2, u, u1, u2, v, v1, v2, x, x1, x2;

	for (o = 0, o_ = 0; o < eo[i]; o++) {
		int j = ej[i][o];

		if (j != p) {
			dfs1(i, j);
			ej[i][o_++] = j;
		}
	}
	eo[i] = o_;
	memset(dq[i], 0x3f, sizeof dq[i]);
	if (cc[i] == 0)
		dq[i][0][1][0] = d + 1;
	else if (cc[i] == 1)
		dq[i][1][1][0] = d + 1;
	else
		dq[i][0][1][0] = dq[i][1][1][0] = d + 1;
	for (o = 0; o < eo[i]; o++) {
		int j = ej[i][o];

		k = o == 0 ? i : ej[i][o - 1];
		memset(dq[j], 0x3f, sizeof dq[j]);
		for (c1 = 0; c1 < 2; c1++)
			for (u1 = 0; u1 < 2; u1++)
				for (v1 = 0; v1 < 2; v1++)
					for (c2 = 0; c2 < 2; c2++)
						for (u2 = 0; u2 < 2; u2++)
							for (v2 = 0; v2 < 2; v2++) {
								x1 = dq[k][c1][u1][v1], x2 = dp[j][c2][u2][v2];
								if (x1 == INF || x2 == INF)
									continue;
								x2++;
								if (v2 && x2 > d)
									continue;
								c = c1;
								if (c1 == c2) {
									u = 0;
									if (!v1 && !v2 || x1 + x2 <= d)
										v = 0, x = min(x1, x2);
									else
										v = 1, x = max(v1 ? x1 : 0, v2 ? x2 : 0);
								} else {
									u = x2 <= d ? 0 : u1;
									if (!v1 && !u2 || x1 + 1 <= d)
										v = 0, x = min(x1, 1);
									else
										v = 1, x = max(v1 ? x1 : 0, u2 ? 1 : 0);
								}
								dq[j][c][u][v] = min(dq[j][c][u][v], x);
							}
	}
	k = eo[i] == 0 ? i : ej[i][eo[i] - 1];
	memcpy(dp[i], dq[k], sizeof dq[k]);
}

void dfs2(int i, int c_, int u_, int v_) {
	int k, o, c, c1, c2, u, u1, u2, v, v1, v2, x, x1, x2;

	memset(dq[i], 0x3f, sizeof dq[i]);
	if (cc[i] == 0)
		dq[i][0][1][0] = d + 1;
	else if (cc[i] == 1)
		dq[i][1][1][0] = d + 1;
	else
		dq[i][0][1][0] = dq[i][1][1][0] = d + 1;
	cc[i] = c_;
	for (o = eo[i] - 1; o >= 0; o--) {
		int j = ej[i][o];

		k = o == 0 ? i : ej[i][o - 1];
		for (c1 = 0; c1 < 2; c1++)
			for (u1 = 0; u1 < 2; u1++)
				for (v1 = 0; v1 < 2; v1++)
					for (c2 = 0; c2 < 2; c2++)
						for (u2 = 0; u2 < 2; u2++)
							for (v2 = 0; v2 < 2; v2++) {
								x1 = dq[k][c1][u1][v1], x2 = dp[j][c2][u2][v2];
								if (x1 == INF || x2 == INF)
									continue;
								x2++;
								if (v2 && x2 > d)
									continue;
								c = c1;
								if (c1 == c2) {
									u = 0;
									if (!v1 && !v2 || x1 + x2 <= d)
										v = 0, x = min(x1, x2);
									else
										v = 1, x = max(v1 ? x1 : 0, v2 ? x2 : 0);
								} else {
									u = x2 <= d ? 0 : u1;
									if (!v1 && !u2 || x1 + 1 <= d)
										v = 0, x = min(x1, 1);
									else
										v = 1, x = max(v1 ? x1 : 0, u2 ? 1 : 0);
								}
								if (c == c_ && u == u_ && v == v_ && x == dq[j][c][u][v]) {
									dfs2(j, c2, u2, v2);
									c_ = c1, u_ = u1, v_ = v1;
									goto out;
								}
							}
out:;
	}
}

int main() {
	int t;

	scanf("%d", &t);
	while (t--) {
		static int dd[N], dd_[N], rr[N], qu[N];
		int cnt, h, i, j, c, o, lower, upper;

		scanf("%d", &n);
		for (i = 0; i < n; i++)
			ej[i] = (int *) malloc(2 * sizeof *ej[i]), eo[i] = 0;
		for (h = 0; h < n - 1; h++) {
			scanf("%d%d", &i, &j), i--, j--;
			append(i, j), append(j, i);
		}
		for (i = 0; i < n; i++)
			scanf("%d", &cc[i]);
		for (i = 0; i < n; i++)
			dd[i] = n;
		cnt = 0;
		for (i = 0; i < n; i++)
			if (cc[i] == -1)
				dd[i] = 0, qu[cnt++] = i;
		for (h = 0; h < cnt; h++) {
			i = qu[h], d = dd[i] + 1;
			for (o = eo[i]; o--; ) {
				j = ej[i][o];
				if (dd[j] > d)
					dd[j] = d, qu[cnt++] = j;
			}
		}
		memcpy(dd_, dd, n * sizeof *dd);
		for (c = 0; c < 2; c++) {
			for (i = 0; i < n; i++)
				dd[i] = n;
			cnt = 0;
			for (i = 0; i < n; i++)
				if (cc[i] == c)
					dd[i] = 0, rr[i] = i, qu[cnt++] = i;
			for (h = 0; h < cnt; h++) {
				i = qu[h], d = dd[i] + 1;
				for (o = eo[i]; o--; ) {
					j = ej[i][o];
					if (dd[j] > d)
						dd[j] = d, rr[j] = rr[i], qu[cnt++] = j;
				}
			}
			for (i = 0; i < n; i++)
				for (o = eo[i]; o--; ) {
					j = ej[i][o];
					if (rr[i] != rr[j])
						dd_[rr[i]] = min(dd_[rr[i]], dd[i] + dd[j] + 1), dd_[rr[j]] = min(dd_[rr[j]], dd[i] + dd[j] + 1);
				}
		}
		lower = 1;
		for (i = 0; i < n; i++)
			lower = max(lower, dd_[i]);
		upper = min(lower + 3, n), lower--;
		while (upper - lower > 1) {
			d = (lower + upper) / 2;
			dfs1(-1, 0);
			if (dp[0][0][0][0] != INF || dp[0][1][0][0] != INF)
				upper = d;
			else
				lower = d;
		}
		if (upper == n)
			printf("-1\n");
		else {
			printf("%d\n", upper);
			d = upper;
			dfs1(-1, 0);
			dfs2(0, dp[0][0][0][0] != INF ? 0 : 1, 0, 0);
			for (i = 0; i < n; i++)
				printf("%d ", cc[i]);
			printf("\n");
		}
		for (i = 0; i < n; i++)
			free(ej[i]);
	}
	return 0;
}

Compilation message (stderr)

balancedtree.c: In function 'append':
balancedtree.c:18:23: warning: suggest parentheses around '-' in operand of '&' [-Wparentheses]
   18 |  if (o >= 2 && (o & o - 1) == 0)
      |                     ~~^~~
balancedtree.c: In function 'dfs1':
balancedtree.c:64:18: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
   64 |          if (!v1 && !v2 || x1 + x2 <= d)
      |              ~~~~^~~~~~
balancedtree.c:70:18: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
   70 |          if (!v1 && !u2 || x1 + 1 <= d)
      |              ~~~~^~~~~~
balancedtree.c: In function 'dfs2':
balancedtree.c:112:18: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  112 |          if (!v1 && !v2 || x1 + x2 <= d)
      |              ~~~~^~~~~~
balancedtree.c:118:18: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  118 |          if (!v1 && !u2 || x1 + 1 <= d)
      |              ~~~~^~~~~~
balancedtree.c: In function 'main':
balancedtree.c:136:2: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
  136 |  scanf("%d", &t);
      |  ^~~~~~~~~~~~~~~
balancedtree.c:141:3: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
  141 |   scanf("%d", &n);
      |   ^~~~~~~~~~~~~~~
balancedtree.c:145:4: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
  145 |    scanf("%d%d", &i, &j), i--, j--;
      |    ^~~~~~~~~~~~~~~~~~~~~
balancedtree.c:149:4: warning: ignoring return value of 'scanf' declared with attribute 'warn_unused_result' [-Wunused-result]
  149 |    scanf("%d", &cc[i]);
      |    ^~~~~~~~~~~~~~~~~~~

• Subtask 1: 10 / 10
• Subtask 2: 13 / 13
• Subtask 3: 26 / 26
• Subtask 4: 45 / 45
• Subtask 5: 6 / 6