Python solution for USACO Silver Question Wormhole Sort

I'm struggling a bit with a USACO silver question using python, .

Basically, I think my code is correct, but I can't figure out a way to make it more efficient as I simply run out of time in the 5 test cases with larger input data.

Here's my code:

import sys sys.setrecursionlimit(200000) #file open fin = open("wormsort.in", "r") fout = open("wormsort.out", "w") #read file temp = fin.readline().strip().split(" ") n, m = int(temp[0]), int(temp[1]) cows = list(map(int, fin.readline().strip().split(" "))) temp = [i for i in range(1, n+1)] #if cows are already sorted if cows == temp: fout.write(str(-1)+ "\n") quit() adjacency = {i: set() for i in range(1, n + 1)} #sorting wormhole by weight weight = [] for i in range(m): temp = list(map(int, fin.readline().strip().split(" "))) #temp storage for wormhold read weight.append(temp[2]) adjacency[temp[1]].add((temp[0], temp[2])) adjacency[temp[0]].add((temp[1], temp[2])) weight.sort() tvis = [0 for _ in range(n)] def dfs(pre, bound, color): #dfs for a single component tvis[pre[0]-1] = color for i in adjacency[pre[0]]: if not tvis[i[0]-1] and i[1] >= bound: dfs(i, bound, color) def check(k): #check if match condition given a min weight k counter = 0 for i in range(1, n+1): counter += 1 if tvis[i-1] == 0: dfs((i, 10**9), k, counter) else: continue for i in range(len(tvis)): if tvis[cows[i]-1] == tvis[i]: continue else: return False return True high = m-1 low = 0 #binary search while low != high: tvis = [0 for _ in range(n)] mid = (high+low)//2 if check(weight[mid]): low = mid+1 else: high = mid fout.write(str(weight[low-1])+ "\n") 

The idea is that since I am trying to find a maximum least weight wormhole, I could use binary search to improve efficiency compared to linear, and in the linear search use dfs for every check to see if each connected component have both the position and the cow included.

2 Answers

If I'm correct about your code running in O(n(n+m)*log n) --- log n is from the binary search, n is from the for loop inside of def check(k), and n+m is from the dfs --- and it still doesn't run under the time limit, there is a possibility that the problem may be impossible to solve in Python.

In the USACO Contest Instructions and Rules, it says, "On some problems (particularly in higher divisions), it may not be possible to solve all inputs within the requisite running time using Python, due to its slowness."

You might have to switch to another language for this problem like C++

You can solve with DSU using python however that is a gold topic