Codeforces Round #303 Div2 E: Paths and Trees (ダイクストラ法)

問題

codeforces.com

解法

まずダイクストラしてスタートからの距離を求めておく。あとは、スタート地点に戻っていくイメージで必要な道を貪欲に集めていけば良い。

コード

import java.util.ArrayList;
import java.util.Arrays;
import java.util.PriorityQueue;
import java.util.Scanner;

public class Main {

	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int N = sc.nextInt();
		if (N == 1) {
			System.out.println(0);
			return;
		}
		int M = sc.nextInt();
		Graph graph = new Graph(N);
		for (int i = 0; i < M; i++) {
			int x = sc.nextInt() - 1;
			int y = sc.nextInt() - 1;
			int cost = sc.nextInt();
			graph.addBidirectionalEdge(i + 1, x, y, cost);
		}

		int start = sc.nextInt() - 1;
		long[] dist = graph.minDistDijkstra(start);
		ArrayList<Integer> ans = new ArrayList<>();

		long sum = 0;
		for (int from = 0; from < N; from++) {
			if (from == start) {
				continue;
			}
			int min = Integer.MAX_VALUE;
			int id = -1;
			for (Edge e : graph.graph[from]) {
				if (dist[from] - e.cost == dist[e.to]) {
					if (min > e.cost) {
						min = e.cost;
						id = e.id;
					}
				}
			}
			sum += min;
			ans.add(id);
		}
		System.out.println(sum);
		StringBuilder builder = new StringBuilder();
		for (Integer integer : ans) {
			builder.append(integer);
			builder.append(" ");
		}
		System.out.println(builder.toString().trim());
		sc.close();

	}
}

class Graph {
	public static final long INF = Long.MAX_VALUE / 2;
	int n;
	ArrayList<Edge>[] graph;

	@SuppressWarnings("unchecked")
	public Graph(int n) {
		this.n = n;
		this.graph = new ArrayList[n];
		for (int i = 0; i < n; i++) {
			graph[i] = new ArrayList<Edge>();
		}
	}

	public void addBidirectionalEdge(int id, int from, int to, int cost) {
		addEdge(id, from, to, cost);
		addEdge(id, to, from, cost);
	}

	public void addEdge(int id, int from, int to, int cost) {
		graph[from].add(new Edge(id, to, cost));
	}

	// dijkstra O(ElogV)
	public long[] minDistDijkstra(int start) {
		long[] dist = new long[n];
		Arrays.fill(dist, INF);
		dist[start] = 0;

		PriorityQueue<Node> priorityQueue = new PriorityQueue<Node>();
		priorityQueue.offer(new Node(0, start));

		while (!priorityQueue.isEmpty()) {
			// キューから1番距離の近いノードを取り出す
			Node node = priorityQueue.poll();
			int v = node.id;
			if (dist[v] < node.dist) {
				// 暫定の最短距離よりも遠かったらスルー
				continue;
			}

			for (Edge e : graph[v]) {
				/*
				 * 取り出したノードから出ている全ての辺について調べ、 暫定の最短距離が更新される場合は更新してキューに入れる
				 */
				if (dist[e.to] > dist[v] + e.cost) {
					dist[e.to] = dist[v] + e.cost;
					priorityQueue.add(new Node(dist[e.to], e.to));
				}
			}
		}
		return dist;
	}

}

class Edge {
	int id;
	int to;
	int cost;

	public Edge(int id, int to, int cost) {
		this.id = id;
		this.to = to;
		this.cost = cost;
	}
}

class Node implements Comparable<Node> {
	long dist;
	int id;

	public Node(long dist, int i) {
		this.dist = dist;
		this.id = i;
	}

	public int compareTo(Node o) {
		if (this.dist > o.dist) {
			return 1;
		} else if (this.dist < o.dist) {
			return -1;
		} else {
			return 0;
		}
	}
}