In [1]:
%%time
fin = open ('15.in', 'r')
fout = open ('clocktree.out', 'w')
N = int(fin.readline().strip())
# use mod 12 of clock times
rawtimes = [int(x)%12 for x in fin.readline().split()]
print(N)
adj = dict() # lookup adjacent nodes to key node
degree = [0]*N # lookup degree of node
# load up adjacency and degree, decrement nodes to start index at 0
for i in range(N-1):
c1,c2 = map(int, fin.readline().split())
c1-=1
c2-=1
degree[c1]+=1
degree[c2]+=1
if c1 in adj:
adj[c1].append(c2)
else:
adj[c1]=[c2]
if c2 in adj:
adj[c2].append(c1)
else:
adj[c2]=[c1]
#print(adj)
#print(degree)
In [2]:
def stepdfs(curr,last1):
global worktimes
# increment clock for node you just stepped into
worktimes[curr]=(worktimes[curr]+1)%12
# loop on all adjacent nodes, but not the node you just came from
for next1 in adj[curr]:
if next1!=last1:
if degree[next1]==1 and worktimes[next1]==0:
continue # ignore leaf nodes that already have correct time
elif degree[next1]==1:
# step repeatedly into leaf nodes to get time correct
delta = 12-worktimes[next1]
worktimes[next1]=(worktimes[next1]+delta)%12
worktimes[curr]=(worktimes[curr]+delta)%12
else:
# recurse down into non-leaf nodes, return last time setting
# so that you can step repeatedly to set it correctly before moving on
delta = 12-stepdfs(next1,curr)
worktimes[curr]=(worktimes[curr]+delta)%12
# update the time on the node you are about to return to
worktimes[last1]=(worktimes[last1]+1)%12
return worktimes[curr]
In [3]:
%%time
count = 0 # count of "good" starting nodes
# loop over all starting nodes
for i in range(N):
# make a copy of original clock times for this loop
worktimes = rawtimes.copy()
# for this starting node, check its immediate neighbors
for j in adj[i]:
if degree[j]==1 and worktimes[j]==0:
continue # ignore leaf nodes that already have correct time
elif degree[j]==1:
# step repeatedly into leaf nodes to get time correct
delta = 12-worktimes[j]
worktimes[j]=(worktimes[j]+delta)%12
worktimes[i]=(worktimes[i]+delta)%12
else:
# recurse down into non-leaf nodes, return last time setting
# so that you can step repeatedly to set it correctly before moving on
delta = 12-stepdfs(j,i)
worktimes[i]=(worktimes[i]+delta)%12
# Stepping back into the starting node is optional
# so if time is set to 1, that's still good because you can ignore that step
if worktimes[i]==0 or worktimes[i]==1:
count+=1
print(count)
fout.write(str(count)+"\n")
fout.close()
This code fails on time for data set 15. It passes all other data sets.¶
I don't believe that it is possible to pass data set 15 with an O(N^2) algorithm in python, but I will keep thinking about it.
An O(N) algorithm is discussed in the official USACO solution, but it is beyond my understanding and beyond the scope of this class.