Submission #261242

#	Time	Username	Problem	Language	Result	Execution time	Memory
261242		idk321	Split the Attractions (IOI19_split)	Java	40 / 100	1603 ms	60880 KiB

split

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

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);
		}

		int[] freq = new int[n];
		for (int i = 0; i < p.length; i++) {
			freq[p[i]]++;
			freq[q[i]]++;
		}
		int maxFreq = 0;
		for (int i = 0; i < n; i++) maxFreq = Math.max(maxFreq, freq[i]);

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

			Arrays.fill(groups, ar[2][0]);
			markNodeA(treeNodeA, treeNodeB, ar[0][0]);
			markNodeB(treeNodeB, treeNodeA, ar[1][0]);
		} else if (maxFreq <= 2) {
			int x = p[0];
			int y = q[0];
			Arrays.fill(groups, 3);
			this.a = a;
			this.b = b;
			markNodeA(x, y, 1);
			markNodeB(y, x, 2);
		} else if (a == 1) {
			this.b = b;
			Arrays.fill(groups, -1);
			markNodes(0);
			for (int i = 0;  i< n; i++) {
				if (groups[i] == 0 || groups[i] == -1) {
					groups[i] = 1;
					break;
				}
			}
			for (int i = 0; i < n; i++) {
				if (groups[i] == 0 || groups[i] == -1) groups[i] = 3;
			}
		}

		return groups;
	}

	private static class ArrayComparator implements Comparator<int[]> {
		@Override
		public int compare(int[] o1, int[] o2) {
			return Integer.compare(o1[1], o2[1]);
		}
	}

	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 markNodes(int node) {
		if (b > 0) {
			groups[node] = 2;
			b--;
		} else groups[node] = 0;
		for (int adjacent : adj[node]) {
			if (groups[adjacent] == -1) {
				markNodes(adjacent);
			}
		}
	}

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

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

	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: 7 / 7
• Subtask 2: 11 / 11
• Subtask 3: 22 / 22
• Subtask 4: 0 / 24
• Subtask 5: 0 / 36