USACO Milk Visits, initial solution

Milk Visits, initial solution¶

USACO 2019 December Contest, Silver¶

http://usaco.org/index.php?page=viewproblem2&cpid=968

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  # Adjacency List

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)

CPU times: user 291 ms, sys: 12.1 ms, total: 303 ms
Wall time: 301 ms

In [2]:

%%time
# loop through all nodes (farms)
# if farm is unvisited (boolean look-up list)
# use BFS to generate a component set for 
# all connected nodes with the same milk flavor 
# and mark all connected farms as visited

listofcompsets = []
visited = [False]*N

def bfs(initial,flavor):
    
    queue = [initial]
    workset = {initial}
    
    while queue:
        
        farm = queue.pop(0)
        if not visited[farm]:
            
            visited[farm] = True
            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 not visited[f]:
        bfs(f,cowtype[f])
            

#print(listofcompsets)            

CPU times: user 205 ms, sys: 11.8 ms, total: 217 ms
Wall time: 216 ms

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.
# This approach is slow, since we loop through all component sets
# until either the start or end farm is found in the set.
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:
        for compset in listofcompsets:
            if (start-1) in compset or (end-1) in compset:
                if (start-1) in compset and (end-1) in compset:
                    satisfied[i]=0
                else:
                    break

CPU times: user 7.47 s, sys: 7.65 ms, total: 7.48 s
Wall time: 7.48 s

In [4]:

ans = ''
for s in satisfied:
    ans += str(s)

fout.write(ans+"\n")
fout.close()