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();
    }
}