Problem Link
Hi, I have written some python code for Wormhole Sort, but it is timing out on anything greater than N=1000. I don’t really know why. My logic is to binary search for the answer k, and check if the k widest wormholes suffice. My check function builds a graph using the first k wormholes and check if every unsorted cow can be reached by traversing that graph. Seems like my code is O(N^2), even though I don’t really know why it is not O(N). My question is are there anything wrong with my code or is python too slow for this problem? Any help will be appreciated, thank you!

import sys
sys.stdin = open("wormsort.in", "r")
sys.stdout = open("wormsort.out", "w")
sys.setrecursionlimit(50000)
n, m = list(map(int, input().split()))
def dfs(graph, node, visited):
if node not in visited:
visited.add(node)
for children in graph[node]:dfs(graph, children, visited)
def check(n):
global wormholes, unsorted_cows
unsorted = unsorted_cows.copy()
firstn = wormholes[:n]
graph = {}
for path in firstn:
a, b, c = path
try: graph[a].append(b)
except: graph[a] = [b]
try: graph[b].append(a)
except: graph[b] = [a]
visited = set()
try:
dfs(graph, unsorted[0], visited)
solvable = True
for cow in unsorted:
if cow not in visited:solvable = False
return solvable
except KeyError:return False
cows = list(map(int, input().split()))
unsorted_cows = []
sorted_cows = cows.copy()
sorted_cows.sort()
for index, cow in enumerate(cows):
if cow != sorted_cows[index]:
unsorted_cows.append(cow)
if unsorted_cows == []:print(-1)
else:
wormholes = []
for i in range(m):
wormholes.append(list(map(int, input().split())))
wormholes.sort(key=lambda x:x[2], reverse=True)
def solve(check, lo, hi):
hi += 1
while lo < hi:
mid = (lo + hi)//2
if check(mid):hi = mid
else:lo = mid+1
return lo
print(wormholes[solve(check, 1, len(wormholes))-1][2])

Ok, so based on your time complexity, it seems like it shouldn’t time out. I think it’s because you’re writing it in Python, which generally times out for higher division problem. I recommend you try rewriting in C++ or Java.