Submission #261204

#	Time	Username	Problem	Language	Result	Execution time	Memory
261204		idk321	Split the Attractions (IOI19_split)	Java	0 / 100	83 ms	10232 KiB

split

import java.util.ArrayList;
import java.util.Arrays;

class split {
	ArrayList<Integer>[] adj;
	int a;
	int b;
	int n;
	int treeNodeA;
	int treeNodeB;
	int[] subtreeSize;
	int[] groups;

	public int[] find_split(int n, int a, int b, int c, int[] p, int[] q) {

		this.n = n;
		groups = new int[n];
		adj = new ArrayList[n];
		for (int i = 0; i < n; i++) adj[i] = new ArrayList<Integer>();
		for (int i = 0; i < p.length; i++) {
			int x = p[i];
			int y = q[i];
			adj[x].add(y);
			adj[y].add(x);
		}

		if (n == p.length + 1) {
			int orgA = a;
			int orgB = b;
			int orgC = c;
			int[] ar = {a, b, c};
			Arrays.sort(ar);
			a = ar[0];
			b = ar[1];
			subtreeSize = new int[n];
			buildTree(0, -1);
			treeNodeA = -1;
			treeNodeB = -1;
			findParts(0, -1);
			if (treeNodeA == -1 || treeNodeB == -1) return groups;

			if (ar[2] == orgA) Arrays.fill(groups, 1);
			else if (ar[2] == orgB) Arrays.fill(groups, 2);
			else Arrays.fill(groups, 3);

			if (orgA == a) {
				markNode(treeNodeA, treeNodeB, 1, a);
			} else if (orgB == a) {
				markNode(treeNodeA, treeNodeB, 2, a);
			} else {
				markNode(treeNodeA, treeNodeB, 3, a);
			}

			if (orgA == b) {
				markNode(treeNodeB, treeNodeA, 1, b);
			} else if (orgB == b) {
				markNode(treeNodeB, treeNodeA, 2, b);
			} else {
				markNode(treeNodeB, treeNodeA, 3, b);
			}
		}

		return groups;
	}

	private void buildTree(int node, int parent) {
		for (int adjacent : adj[node]) {
			if (adjacent == parent) continue;
			buildTree(adjacent, node);
			subtreeSize[node] += subtreeSize[adjacent];
		}
		subtreeSize[node]++;
	}

	private void markNode(int node, int parent, int group, int needed) {
		if (needed == 0) return;
		groups[node] = group;
		needed--;
		for (int adjacent : adj[node]) {
			if (adjacent == parent) continue;
			markNode(adjacent, node, group, needed);
		}
	}

	private void findParts(int node, int parent) {
		for (int adjacent : adj[node]) {
			if (adjacent == parent) continue;
			if (subtreeSize[adjacent] >= a) {
				if (n - subtreeSize[adjacent] >= b) {
					treeNodeA = adjacent;
					treeNodeB = node;
				}
			}
			if (subtreeSize[adjacent] >= b) {
				if (n - subtreeSize[adjacent] >= a) {
					treeNodeA = node;
					treeNodeB = adjacent;
				}
			}
			findParts(adjacent, node);
		}
	}
}

Compilation message (stderr)

Note: split.java uses unchecked or unsafe operations.
Note: Recompile with -Xlint:unchecked for details.

• Subtask 1: 0 / 7
• Subtask 2: 0 / 11
• Subtask 3: 0 / 22
• Subtask 4: 0 / 24
• Subtask 5: 0 / 36