[leetcode 周赛] 100276. 最短路径中的边

solution

• 使用dijkstra算法求出顶点0到每个顶点的最小距离dp[i]
• 然后再从n-1开始遍历，如果dp[to] + w == dp[from] ，这条边符合条件

import java.util.*;

class Solution {
    public static void main(String[] args) {
        Solution solution = new Solution();
        boolean[] ans = solution.findAnswer(6, new int[][] {
                {0,1,4},
                {0,2,1},
                {1,3,2},{1,4,3},{1,5,1},{2,3,1},{3,5,3},{4,5,2}
        });
    }
    public boolean[] findAnswer(int n, int[][] edges) {
        List<int[]>[] g = new ArrayList[n];
        for (int i = 0; i < n; i++) {
            g[i] = new ArrayList<>();
        }
        int[] dp = new int[n];
        Arrays.fill(dp,Integer.MAX_VALUE);

        Map<Integer,Map<Integer,Integer>> map = new HashMap<>();
        int index = 0;
        for (int[] edge : edges) {
            int from = edge[0];
            int to = edge[1];
            int w = edge[2];
            g[from].add(new int[]{to, w});
            g[to].add(new int[]{from, w});

            map.put(from, map.getOrDefault(from,new HashMap<>()));
            map.put(to, map.getOrDefault(to, new HashMap<>()));
            map.get(from).put(to, index);
            map.get(to).put(from, index);
            index++;
        }
        //bfs;
        dp[0] = 0;
        PriorityQueue<Integer> queue = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return dp[o1] - dp[02];
            }
        });
        queue.add(0);

        while (!queue.isEmpty()) {
            int from = queue.poll();
//            visited[from] = true;
            for (int[] edge: g[from]) {
                int to = edge[0];
                int w = edge[1];

                if(dp[to] > dp[from] + w) {
                    dp[to] = dp[from] + w;
                    queue.add(to);
                }
            }
        }
        System.out.println(dp[n-1]);
        //
        boolean[] answer = new boolean[edges.length];
        int curr = n - 1;
        ArrayDeque<Integer> arrayDeque = new ArrayDeque<>();
        boolean[] visited = new boolean[n];
//        visited[curr] = true;
        arrayDeque.add(curr);
        while (!arrayDeque.isEmpty()) {
          int from = arrayDeque.poll();
          for (int[] edge: g[from]) {
              int to = edge[0];
              int w = edge[1];

              if (dp[to] == dp[from] - w) {
//                  visited[to] = true;
                  int idx = map.get(from).get(to);
                  answer[idx] = true;
                  arrayDeque.add(to);
              }
          }
        }
        return answer;
    }
}