Wormhole Sort TLE

Problems and Contests
#1

Hi,

I used Java to code a solution for Wormhole Sort. The solution binary searches on the least wide wormhole, and constructs a graph by drawing an edge between two cows connected by a wormhole. Then, using DFS and HashSet, it calculates if all components are valid in linear time.

However, this solution passes the first nine test cases, but gets TLE on the tenth one. Is this because my code is poorly optimized, or Java is too slow?

Here is my code:

import java.io.*;
import java.util.*;

public class WormholeSort {
    public static void main(String[] args) throws IOException {
        // BufferedReader in = new BufferedReader(new InputStreamReader(System.in));
        // PrintWriter out = new PrintWriter(System.out);
       BufferedReader in = new BufferedReader(new FileReader("wormsort.in"));
       PrintWriter out = new PrintWriter(new FileWriter("wormsort.out"));
        StringTokenizer st = new StringTokenizer(in.readLine());
        int numCows = Integer.parseInt(st.nextToken());
        int numWormholes = Integer.parseInt(st.nextToken());
        int[] cows = new int[numCows];
        StringTokenizer stc = new StringTokenizer(in.readLine());
        for (int i = 0; i < numCows; i++) {
            cows[i] = Integer.parseInt(stc.nextToken()) - 1;
        }
        Wormhole[] wormholes = new Wormhole[numWormholes];
        for (int i = 0; i < numWormholes; i++) {
            StringTokenizer stw = new StringTokenizer(in.readLine());
            wormholes[i] = new Wormhole(Integer.parseInt(stw.nextToken()) - 1,
                Integer.parseInt(stw.nextToken()) - 1, Integer.parseInt(stw.nextToken()));
        }
        // sort by width
        Arrays.sort(wormholes);
        // Binary Search on least wide wormhole
        int low = 0;
        int high = numWormholes - 1;
        while (low <= high) {
            int mid = (low + high) / 2;
            boolean midCanSort = canSort(numCows, cows, wormholes, mid);
            boolean higherCanSort = canSort(numCows, cows, wormholes, mid + 1);
            if (midCanSort && !higherCanSort) {
                out.println(wormholes[mid].width);
                out.close();
                return;
            } else if (midCanSort) {
                low = mid + 1;
            } else {
                high = mid - 1;
            }
        }
        out.println(-1);
        out.close();
    }
    static boolean canSort(int numCows, int[] cows, Wormhole[] wormholes, int minHole) {
        // Indices are treated as node, wormholes as edges
        ArrayList<ArrayList<Integer>> adjacent = new ArrayList<>(numCows);
        for (int i = 0; i < numCows; i++) {
            adjacent.add(new ArrayList<>());
        }
        for (int i = minHole; i < wormholes.length; i++) {
            adjacent.get(wormholes[i].p1).add(wormholes[i].p2);
            adjacent.get(wormholes[i].p2).add(wormholes[i].p1);
        }
        // Boolean array to check if index's connected component is already checked
        boolean[] checked = new boolean[numCows];
        for (int i = 0; i < numCows; i++) {
            if (!checked[i]) {
                // ArrayList with all indices in connected component
                ArrayList<Integer> indices = new ArrayList<>();
                floodFill(adjacent, checked, indices, i);
                // HashSet with all cows in those indices
                HashSet<Integer> cowsInIndices = new HashSet<>();
                for (int index : indices) {
                    cowsInIndices.add(cows[index]);
                }
                for (int index : indices) {
                    if (!cowsInIndices.contains(index)) {
                        return false;
                    }
                }
            }
        }
        return true;
    }
    static void floodFill(ArrayList<ArrayList<Integer>> adjacent, boolean[] checked, ArrayList<Integer> indices, int curIndex) {
        if (checked[curIndex]) {
            return;
        }
        checked[curIndex] = true;
        indices.add(curIndex);
        for (int adj : adjacent.get(curIndex)) {
            floodFill(adjacent, checked, indices, adj);
        }
    }
}
class Wormhole implements Comparable<Wormhole> {
    public int p1;
    public int p2;
    public int width;

    public Wormhole(int p1, int p2, int width) {
        this.p1 = p1;
        this.p2 = p2;
        this.width = width;
    }

    public int compareTo(Wormhole wh) {
        return this.width - wh.width;
    }
}

Thanks,
rayfish

#2

It seems your solution uses recursion. While that’s fine, it can result in a large constant factor for your code. Have you tried an iterative approach instead?

1 Like
#3

You don’t have to check both midCanSort and higherCanSort on each binary search iteration (which doubles your time complexity). You can instead set low = 0 and high = numWormholes and find the lastTrue index (highest mid that works). If this lastTrue is equal to numWormholes, then no wormholes are necessary to sort the cows, so you can print -1. Else, print the answer wormholes[ans].width – on each iteration that mid works, you’ll have to set ans = mid.

1 Like