USACO Silver "Wormhole Sort" solver

Question

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.

Answer 1

Layout

It is common to place functions after the import lines and before executable code. Having them in the middle of the code interrupts the natural flow of the code (from a human readability standpoint).

Documentation

The PEP 8 style guide recommends adding docstrings for the main code and functions. For example, you could convert the function comment into a docstring:

def check(k):
    """ check if match condition given a min weight k """

Also, you could describe the purpose of the code in a docstring at the top of the file. This would include a description of the wormhole solver question.

Simpler

These 2 lines can be deleted from the check function:

    else:
        continue

Similarly, this if/else:

    if tvis[cows[i]-1] == tvis[i]:
        continue
    else:
        return False

can be simplified as:

    if tvis[cows[i]-1] != tvis[i]:
        return False

Large integer values like:

200000

are easier to read using underscores:

200_000