[ABC252] E - Road Reduction (AtCoder Beginner Contest 252 in Java)

문제 : abc252_e

 

  다익스트라 자료구조에 대해 이해를 하고있다면 쉽게 풀 수 있다. 문제에서 제시된 d2+d3+...+dN의 합을 최소화 한다는 말은 즉, 1부터 모든 정점으로의 최단거리를 가지는 N-1개의 edge를 남겨야 한다는 말과 동일하다. 1부터 모든 정점으로의 최단거리는 다익스트라를 돌리면 얻을 수 있고, 다익스트라의 결과는 모든 정점으로의 경로가 존재했다면 N-1개만 남게 된다.

 

  그러므로, 단순히 정점 1부터 시작하는 다익스트라를 구현하고, 이 때 각 정점으로 마지막으로 들어왔던 간선(즉, 각 정점으로 들어온 1부터 최단거리에 해당하는 간선)을 저장해둔 후 그걸 출력해주면 된다. 이 때 'in arbitrary order' 라고 했으므로(랜덤 순서를 뜻한다) 아무 순서로든 출력만 해주면 된다.

 

  이하 코드에서는 selectedEdgeIdx가 위에서 말한 각 정점에 마지막으로 들어왔던 간선을 뜻한다.

 

코드 : github

...
class Edge {
    int to, dist, idx;
    public Edge(int to, int dist, int idx) {
        this.to = to;
        this.dist = dist;
        this.idx = idx;
    }
}
class Pos implements Comparable<Pos> {
    int num;
    long dist;
    public Pos(int num, long dist) {
        this.num = num;
        this.dist = dist;
    }

    @Override
    public int compareTo(Pos o) {
        long calc = this.dist - o.dist;
        if (calc < 0) return -1;
        else if (calc > 0) return 1;
        return 0;
    }
}
int n;
ArrayList<Edge>[] edges;

private int[] dijkstra() {
    long[] dist = new long[n+1];
    Arrays.fill(dist, Long.MAX_VALUE);
    int[] selectedEdgeIdx = new int[n+1];
    PriorityQueue<Pos> pq = new PriorityQueue<>();
    pq.add(new Pos(1, 0));
    dist[0] = 0;
    while (!pq.isEmpty()) {
        Pos cur = pq.poll();
        if (cur.dist > dist[cur.num]) continue;
        for (Edge edge : edges[cur.num]) {
            int next = edge.to;
            if (dist[next] <= edge.dist + cur.dist) continue;
            dist[next] = cur.dist + edge.dist;
            selectedEdgeIdx[next] = edge.idx;
            pq.add(new Pos(next, dist[next]));
        }
    }
    return selectedEdgeIdx;
}

private void solution() throws Exception {
    n = nextInt();
    int m = nextInt();
    edges = new ArrayList[n+1];
    for (int i = 1; i <= n; i++) edges[i] = new ArrayList<>();
    for (int i = 1; i <= m; i++) {
        int a = nextInt();
        int b = nextInt();
        int c = nextInt();
        edges[a].add(new Edge(b, c, i));
        edges[b].add(new Edge(a, c, i));
    }
    int[] selectedEdgeIdx = dijkstra();
    StringBuilder sb = new StringBuilder();
    for (int i = 2; i <= n; i++) {
        sb.append(selectedEdgeIdx[i]).append(' ');
    }
    System.out.println(sb);
}
...