10 Wizards, Java - rextester

Language: Layout:

//'main' method must be in a class 'Rextester'.
//Compiler version 1.8.0_111

import java.util.*;
import java.lang.*;

class Rextester
{  
    public static void main(String args[])
    {
        List<List<Integer>> wizards = new ArrayList<>();
        wizards.add(Arrays.asList(1, 2));
        wizards.add(Arrays.asList(3, 4, 2));
        wizards.add(Arrays.asList(3, 4));
        wizards.add(Arrays.asList(4));
        wizards.add(Arrays.asList());
        System.out.println(wizards);
        System.out.println(shortestPath(wizards, 0, wizards.size() - 1));
        System.out.println(getShortestPath(wizards, 0, wizards.size() - 1));
    }
    
    public static List<Integer> shortestPath(List<List<Integer>> wizards, int src, int dst) {
        int n = wizards.size();
        int[] dist = new int[n];
        int[] pred = new int[n];
        Arrays.fill(dist, Integer.MAX_VALUE);
        Arrays.fill(pred, -1);
        dist[src] = 0;
        TreeSet<Integer> set = new TreeSet<>((a, b) -> dist[a] - dist[b]);
        for (int i = 0; i < n; i++) set.add(i);
        while (!set.isEmpty()) {
            int u = set.pollFirst();
            for (int v: wizards.get(u)) {
                int w = (u - v)*(u - v);
                if (dist[u] + w < dist[v]) {
                    set.remove(v);
                    dist[v] = dist[u] + w;
                    pred[v] = u;
                    set.add(v);
                }
            }  
        }
        LinkedList<Integer> path = new LinkedList<>();
        if (dist[dst] != Integer.MAX_VALUE) for (int i = dst; i != -1; i = pred[i]) path.addFirst(i);
        return path;
    }
}

[ + ] Show input
Compilation time: 1.34 sec, absolute running time: 0.33 sec, cpu time: 0.37 sec, memory peak: 24 Mb, absolute service time: 1,69 sec