Wormhole Sort Timing Out

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.

That is what I suspected, I am going to rewrite the program in c++. Thank you very much for your help!

Hi, I translated the code above into c++ but it still times out whenever N is greater than 1000. Can someone help? Thank you!


using namespace std;

int n, m;
set<int> visited;
vector<vector<int>> wormholes;
vector<int> unsorted_cows;
vector<int> cows;
vector<int> sorted_cows;

bool cmp(const vector<int>& a, const vector<int>& b) {
  return a[2] > b[2];
}

void dfs(map<int, vector<int>> graph, int node) {
  if (visited.find(node) == visited.end()) {
    visited.insert(node);
    vector<int> children = graph[node];
    for (int i = 0; i < children.size(); i++) {
      dfs(graph, children[i]);
    }
  }
}

bool check(int n) {
  map<int, vector<int>> graph;
  for (int i = 0; i < n; i++) {
    vector<int> path = wormholes[i];
    int a = path[0];
    int b = path[1];
    graph[a].push_back(b);
    graph[b].push_back(a);
  }
  try {
    dfs(graph, unsorted_cows[0]);
    bool solvable = true;
    for (int i = 0; i < unsorted_cows.size(); i++) {
      int cow = unsorted_cows[i];
      if (visited.find(cow) == visited.end()) solvable = false;
    }
    visited.clear();
    return solvable;
  } catch (...) {visited.clear();return false;}
}

int solve(function<bool(int)> f, int lo, int hi) {
    for (hi ++; lo < hi; ) {
        int mid = (lo+hi)/2;
        f(mid) ? hi = mid : lo = mid+1;
    }
    return lo;
}

int main() {
  freopen("wormsort.in", "r", stdin);
  freopen("wormsort.out", "w", stdout);
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cin >> n >> m;
  for (int i = 0; i < n; i++) {
    int a;
    cin >> a;
    cows.push_back(a);
  }
  sorted_cows = cows;
  sort(sorted_cows.begin(), sorted_cows.end());
  for (int i = 0; i < n; i++) {
    if (sorted_cows[i] != cows[i]) {
      unsorted_cows.push_back(cows[i]);
    }
  }
  if (unsorted_cows.size() == 0) {
    cout << -1;
  } else{
      for (int i = 0; i < m; i++) {
        int a, b, c;
        cin >> a >> b >> c;
        wormholes.push_back({a, b, c});
      }
      sort(wormholes.begin(), wormholes.end(), cmp);
      cout << wormholes[solve(check, 1, wormholes.size())-1][2];
  }
  return 0;
}

pass the graph as reference instead

void dfs(map<int, vector<int>> graph, int node) {

should be

void dfs(map<int, vector<int>> &graph, int node) {

or just use global variables

Also, using a map instead of an array of vectors adds an additional log factor to the complexity.

I don’t think this try-catch has any effect.

1 Like