Problem Link
2020 January Silver Problem #3 - Clock Tree - USACO
My Work
import java.io.*;
import java.util.*;
public class clocktree {
public static TreeMap<Integer, TreeSet<Integer>> connections;
public static int[] clocks;
public static void main (String[] args) throws IOException {
// read in data
BufferedReader input = new BufferedReader(new InputStreamReader(new FileInputStream("clocktree.in")));
int N = Integer.parseInt(input.readLine());
clocks = new int[N+1];
StringTokenizer st = new StringTokenizer(input.readLine());
connections = new TreeMap<>();
for (int i=1;i<=N;i++) {
clocks[i] = Integer.parseInt(st.nextToken());
connections.put(i, new TreeSet<Integer>()); }
for (int i=0;i<N-1;i++) {
st = new StringTokenizer(input.readLine());
int a = Integer.parseInt(st.nextToken());
int b = Integer.parseInt(st.nextToken());
connections.get(a).add(b);
connections.get(b).add(a); }
// check if each starting room is valid
// System.out.println(connections);
int count = 0;
for (int i=1;i<=N;i++) if (isPossible(i)) count ++;
// write result
//PrintWriter out = new PrintWriter(System.out);
PrintWriter out = new PrintWriter(new FileWriter("clocktree.out"));
out.println(count);
out.close();
}
static boolean isPossible (int start) {
// check if position is valid starting point
// set up
int n = clocks.length;
int[] clockPos = clocks.clone();
clockPos[0] = 12;
int[] checker = new int[n];
Arrays.fill(checker, 12);
boolean[] changedTo12 = new boolean[n];
LinkedList<int[]> stack = new LinkedList<>();
stack.push(new int[] {start, -1});
// simulate a walk through the tree with a depth-first search
// (where leaves are moved to 12 and removed, and then repeat)
while (!stack.isEmpty()) {
// System.out.println(stack);
int[] s = stack.pop();
int node = s[0];
int parent = s[1];
if (Arrays.equals(clockPos, checker)) return true; // if all clocks at 12 stop and return true
if (parent != -1) clockPos[node] = add(clockPos[node], 1); // move dial forward 1 to represent walking into room
boolean isLeaf = true;
for (int c: connections.get(node)) { // check if node is not a leaf
if (c!=parent && !changedTo12[c]) {
isLeaf = false;
stack.push(new int[] {c, node});
break; } }
if (isLeaf) { // otherwise simulate going back and forth between leaf and parent until leaf is 12 (ending on leaf)
if (!changedTo12[parent]) {
clockPos[parent] = add(clockPos[parent], 12-clockPos[node]);
clockPos[node] = 12;
changedTo12[node] = true;
stack.push(new int[] {parent, node}); // return to parent to continue walk
} }
}
return Arrays.equals(clockPos, checker);
}
static int add (int a, int b) {
// adds b hours to time a on a 12-hour system
while (a+b > 12) {
b -= 12-a;
a = 0; }
return a+b; }
}
This basically solves the problem with an O(N^2) solution. The algorithm starts at a given point on the tree, going all the way to each leaf and setting it to 12 by going back and forth between the leaf and its parent, and then ‘removing’ that leaf from consideration and repeating the process with newly-created leaves. This is done for each of the N nodes, with the number of successful ones counted and output as the answer.
Question
Right now this code passes 12 out of 15 test cases, and times out on the last 3. I think there are probably just slow methods or sections that I haven’t recognized, but I’m just wondering how could I optimize this code to get those last three test cases? I know that there is an O(N) solution, however I’d like to try to get my O(N^2) solution to pass (so I don’t want to change the logic itself too much).