CODE FESTIVAL 2014 Middle A: 身体バランス (ダイクストラ法)

解法

ダイクストラ

コード

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.PriorityQueue;

public class Main {

	public static void main(String[] args) {
		int N = nextInt();
		int M = nextInt();
		int s = nextInt() - 1;
		int t = nextInt() - 1;
		if (s == t) {
			System.out.println(s + 1);
			return;
		}

		Graph graph = new Graph(N);

		for (int i = 0; i < M; i++) {
			int x = nextInt() - 1;
			int y = nextInt() - 1;
			int d = nextInt();
			graph.addBidirectionalEdge(x, y, d);
		}

		int[] sD = graph.minDistDijkstra(s);
		int[] tD = graph.minDistDijkstra(t);
		for (int i = 0; i < N; i++) {
			if (sD[i] == tD[i] && sD[i] != 0 && sD[i] <= 1000) {
				System.out.println(i + 1);
				return;
			}
		}
		System.out.println(-1);
		return;

	}

	static int nextInt() {
		int c;
		try {
			c = System.in.read();
			while (c != '-' && (c < '0' || c > '9'))
				c = System.in.read();
			if (c == '-')
				return -nextInt();
			int res = 0;
			while (c >= '0' && c <= '9') {
				res = res * 10 + c - '0';
				c = System.in.read();
			}
			return res;
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		return -1;
	}

}

class Graph {
	public static final int INF = 1 << 29;
	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 from, int to, int cost) {
		addEdge(from, to, cost);
		addEdge(to, from, cost);
	}

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

	// dijkstra O(ElogV)
	public int[] minDistDijkstra(int start) {
		int[] dist = new int[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 to;
		int cost;

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

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

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

		public int compareTo(Node o) {
			return this.dist - o.dist;
		}
	}
}