I'm struggling a bit with a USACO silver question using Python: http://usaco.org/index.php?page=viewproblem2&cpid=992.
The question provides an unsorted list of numbers (cows) and a number of edges (wormholes), each with a weight, connecting two numbers. The question asks for the maximum weight for which, each edge of that weight or above can be used, where the list can be sorted using only those edges.
Basically, I think my code is correct, as I get the problems with less input size 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 has both the position and the cow included.
