In [1]:
%%time
# read in data and create an adjacency list
# also create a lookup list for cow type ('G' or 'H')
# data set 12 is a big one
fin = open ('12.in', 'r')
fout = open ('milkvisits.out', 'w')
N,M = map(int, fin.readline().split())
# N Farms
# M Friends
cowtype = fin.readline().strip()
empty_list = []
ADJ = [empty_list]*N
for i in range(N-1): # Tree has N-1 edges
x,y = map(int, fin.readline().split())
#print(x,y)
ADJ[x-1]=ADJ[x-1]+[y-1]
ADJ[y-1]=ADJ[y-1]+[x-1]
#print(ADJ)
#print(cowtype)
In [2]:
%%time
# loop through all nodes (farms)
# if farm is unvisited (integer look-up list initalized to -1)
# use BFS to generate a component set for
# all connected nodes with the same milk flavor
# and mark all connected farms as visited by
# updating the visited list with the initial farm index that started
# the BFS search. That farm index is the "parent" of the component set.
listofcompsets = []
visited = [-1]*N
def bfs(initial,flavor):
queue = [initial]
workset = {initial}
while queue:
farm = queue.pop(0)
if visited[farm]==-1:
visited[farm] = initial
neighbors = ADJ[farm]
for neighbor in neighbors:
if cowtype[neighbor]==flavor:
queue.append(neighbor)
workset.add(neighbor)
#print(workset, flavor)
listofcompsets.append(workset)
for f in range(N):
if visited[f]==-1:
bfs(f,cowtype[f])
#print(listofcompsets)
In [3]:
%%time
# If start and end farms are in the same connected component
# then the friend will not be satisfied if the starting farm
# milk type is not the desired choice
# If the start and end farms are not in the same connected component
# then both milk types will be encountered and the friend will be
# satisfied regardless of desired milk type.
# Now it is only necessary to do a fast lookup in the visited list
# to see if the start and end farm have the same parent
# It is not even necessary to store the component sets
satisfied = [1]*M
for i in range(M):
inlist = fin.readline().split()
#print(inlist)
start = int(inlist[0])
end = int(inlist[1])
milk = inlist[2]
#print(start,end,milk)
if cowtype[start-1]!=milk:
if visited[start-1]==visited[end-1]:
satisfied[i]=0
ans = ''
for s in satisfied:
ans += str(s)
fout.write(ans+"\n")
fout.close()
