apjcoder123's blog

Need help with implementation of java code for Dijkstra's algorithm
By apjcoder123, history, 14 months ago, In English

Hi community, I've come across this problem: https://cses.fi/problemset/task/1671/

Below is my Java solution

Issue: This code is resulting into a TLE for some tests, Can anyone please help me out with the implementation issue ?

import java.util.*;

public class Graph2 {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);

        int n = sc.nextInt(); // number of nodes
        int m = sc.nextInt(); // number of edges

        ArrayList<int[]>[] adj = new ArrayList[n];
        for (int i = 0; i < n; i++) adj[i] = new ArrayList<>();

        for (int i = 0; i < m; i++) {
            int u = sc.nextInt() - 1;
            int v = sc.nextInt() - 1;
            int w = sc.nextInt();
            adj[u].add(new int[]{v, w});
        }

        long[] dist = new long[n];
        Arrays.fill(dist, Long.MAX_VALUE);
        dist[0] = 0;

        TreeSet<List<Integer>> set = new TreeSet<>((a, b) -> {
            if (!a.get(0).equals(b.get(0))) return Long.compare(a.get(0), b.get(0));
            return Integer.compare(a.get(1), b.get(1));
        });

        set.add(Arrays.asList(0, 0)); // [distance, node]

        while (!set.isEmpty()) {
            List<Integer> top = set.pollFirst();
            int u = top.get(1);

            for (int[] edge : adj[u]) {
                int v = edge[0], w = edge[1];
                long newDist = dist[u] + w;
                if (newDist < dist[v]) {
                    set.remove(Arrays.asList((int) dist[v], v));
                    dist[v] = newDist;
                    set.add(Arrays.asList((int) newDist, v));
                }
            }
        }

        for (int i = 0; i < n; i++) {
            System.out.print(dist[i] + " ");
        }
        System.out.println();
    }
}
  • Vote: I like it
  • -6
  • Vote: I do not like it

Comments Comments (2)
| Write comment?
»
14 months ago, hide # |
 
Vote: I like it 0 Vote: I do not like it

In place of TreeSet you may use Queue and visited array to remove logn time complexity.

Still it is possible that you get TLE because some CSES problems could be solved using c++ only. Some months back, I had tried two CSES problems to solve in java but each time I got TLE and when tried in C++ with the same logic, it got accepted. I used to believe that it is not possible that some problem could be solved in other language and not in java, until I had read somewhere that if choosing slow algorithm should be punished then why not punish slow language? I would advise you to shift to c++, if you really wanna pursue cp.

→ Reply
»
14 months ago, hide # |
← Rev. 2  
Vote: I like it 0 Vote: I do not like it 
import java.util.*;

public class Graph2 {
    public static void main(String[] args) {
        Scanner sc = new Scanner(System.in);
        int n = sc.nextInt();
        int m = sc.nextInt(); 
        ArrayList<int[]>[] adj = new ArrayList[n];
        for (int i = 0; i < n; i++) {
            adj[i] = new ArrayList<>();
        }
        for (int i = 0; i < m; i++) {
            int u = sc.nextInt() - 1;
            int v = sc.nextInt() - 1;
            int w = sc.nextInt();
            adj[u].add(new int[]{v, w});
        }
        long[] dist = new long[n];
        Arrays.fill(dist, Long.MAX_VALUE);
        dist[0] = 0;
        PriorityQueue<long[]> pq = new PriorityQueue<>((a, b) -> {
            if (a[0] != b[0]) return Long.compare(a[0], b[0]);
            return Long.compare(a[1], b[1]);
        });
        pq.offer(new long[]{0, 0});
        while (!pq.isEmpty()) {
            long[] current = pq.poll();
            long d = current[0];
            int u = (int) current[1];
            if (d != dist[u]) continue;

            for (int[] edge : adj[u]) {
                int v = edge[0];
                int w = edge[1];
                long newDist = d + w;
                if (newDist < dist[v]) {
                    dist[v] = newDist;
                    pq.offer(new long[]{newDist, v});
                }
            }
        }
        for (int i = 0; i < n; i++) {
            System.out.print(dist[i] + " ");
        }
        System.out.println();
    }
}

i think that will be good enough= just use priority queue instead of TreeSet. Also its probably because of some cses is only solvable with c++ as someone said already.

→ Reply