We use cookies to ensure you have the best browsing experience on our website. Please read our cookie policy for more information about how we use cookies.
importmathimportosimportheapq#Importheapqforthepriorityqueueimportreimportsys## Complete the 'shop' function below.## The function is expected to return an INTEGER.# The function accepts following parameters:# 1. INTEGER n# 2. INTEGER k# 3. STRING_ARRAY centers# 4. 2D_INTEGER_ARRAY roads#defshop(n,k,centers,roads):# Step 1: Pre-process the centers input to store fish masksfish_masks=[0]*(n+1)fori,center_fish_strinenumerate(centers):# The input is a string of space-separated integerscenter_fish=list(map(int,center_fish_str.split()))ifnotcenter_fish:continue# The first integer is the count, the rest are fish typesnum_fish=center_fish[0]forjinrange(1,num_fish+1):fish_type=center_fish[j]fish_masks[i+1]|=(1<<(fish_type-1))# Step 2: Build the adjacency list for the graphadj=[[]for_inrange(n+1)]forroadinroads:u,v,w=roadadj[u].append((v,w))adj[v].append((u,w))# Step 3: Run Dijkstra's with a state-space (node, fish_mask)# distances[node][mask] stores the minimum time to reach 'node' having collected fish in 'mask'distances=[[float('inf')]*(1<<k)for_inrange(n+1)]# Priority queue: (time, node, fish_mask)pq=[(0,1,fish_masks[1])]distances[1][fish_masks[1]]=0whilepq:time,u,mask=heapq.heappop(pq)iftime>distances[u][mask]:continueforv,weightinadj[u]:new_mask=mask|fish_masks[v]new_time=time+weightifnew_time<distances[v][new_mask]:distances[v][new_mask]=new_timeheapq.heappush(pq,(new_time,v,new_mask))# Step 4: Calculate the minimum time for two catsmin_total_time=float('inf')full_mask=(1<<k)-1formask1inrange(1<<k):formask2inrange(mask1,1<<k):if(mask1|mask2)==full_mask:# Time is the maximum of the two cats' path timestime_cat1=distances[n][mask1]time_cat2=distances[n][mask2]# Check for unreachable statesiftime_cat1!=float('inf')andtime_cat2!=float('inf'):min_total_time=min(min_total_time,max(time_cat1,time_cat2))returnmin_total_timeif__name__=='__main__':fptr=open(os.environ['OUTPUT_PATH'],'w')first_multiple_input=input().rstrip().split()n=int(first_multiple_input[0])m=int(first_multiple_input[1])k=int(first_multiple_input[2])centers=[]for_inrange(n):centers_item=input()centers.append(centers_item)roads=[]for_inrange(m):roads.append(list(map(int,input().rstrip().split())))res=shop(n,k,centers,roads)fptr.write(str(res)+'\n')fptr.close()
Cookie support is required to access HackerRank
Seems like cookies are disabled on this browser, please enable them to open this website
Synchronous Shopping
You are viewing a single comment's thread. Return to all comments →