I♥TLE

Java、オンラインジャッジなど

今更だけどD問題を解いてみた

だいぶ前に参加したコンテストの問題ですが、しばらく解説を見ても作れないでいたのがようやくACになったので載せておきます。

解法の概要

公式の解答にあるように、街の座標をxまたはyについてソートした場合に隣同士になった街の間にしかルートは必要ないと考えて大丈夫です。
街を頂点、ルートを辺とみなしたグラフの最小全域木のコストを計算すれば良いのですが、行列を使ってグラフを表現するとそもそも速度どころかメモリの確保ができないため、それぞれの街ごとにマップを使ってルートを保存しています。
優先度付きキューを用いたプリム法により時間内で処理を行うことができます(先頭要素を取り出す度に距離が最短の要素が自動的に先頭に来るため)。

public class Main {
 
	public static void main(String[] args) throws NumberFormatException, IOException {

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
 
		int n = Integer.parseInt(br.readLine());
		
		Town[] towns = new Town[n];
		Town[] copy = new Town[n];
		for(int i = 0; i < n; i++){
			String[] tmpArray = br.readLine().split(" ");
			towns[i] = new Town(Integer.parseInt(tmpArray[0]), Integer.parseInt(tmpArray[1]));
			copy[i] = towns[i];
		}
		
		//xについてソート
		Arrays.sort(copy, new Comparator<Town>() {
			public int compare(Town t1, Town t2){
				return t1.x - t2.x;
			}
		});
		
		//隣接する街のコスト計算
		for(int i = 0; i < n - 1; i++){
			int cost = Math.min(copy[i+1].x - copy[i].x, Math.abs(copy[i+1].y - copy[i].y));
			copy[i].addCost(copy[i+1], cost);
			copy[i+1].addCost(copy[i], cost);
		}
		
		//yについてソート
		Arrays.sort(copy, new Comparator<Town>() {
			public int compare(Town t1, Town t2){
				return t1.y - t2.y;
			}
		});
 
		//隣接する街のコスト計算
		for(int i = 0; i < n - 1; i++){
			int cost = Math.min(copy[i+1].y - copy[i].y, Math.abs(copy[i+1].x - copy[i].x));
			copy[i].addCost(copy[i+1], cost);
			copy[i+1].addCost(copy[i], cost);
		}
		
		long result = MST(towns);
		System.out.println(result);
	}
	//プリム法によりコストの合計を求める
	static long MST(Town[] towns){
		long cost = 0;
		int n = towns.length;
		boolean[] visited = new boolean[n];
		
		PriorityQueue<Edge> pq = new PriorityQueue<Edge>();
		//0番目の街から始める
		visited[0] = true;
		Edge[] edges = towns[0].getEdges();
		for(int i = 0; i < edges.length ; i++){
			pq.add(edges[i]);
		}
		
		while(!pq.isEmpty()){
			//最小の辺を探す
			Edge minEdge = pq.poll();

			if(visited[minEdge.to]){
				continue;
			}
			visited[minEdge.to] = true;
			cost += minEdge.cost;
			Town nextTown = towns[minEdge.to];
			
			edges = nextTown.getEdges();
			
			for(int i = 0; i < edges.length ; i++){
				pq.add(edges[i]);
			}			
		}
		
		return cost;
	}
 
}
 
class Edge implements Comparable<Edge>{
	int from;
	int to;
	int cost;
	
	public Edge(Town t1, Town t2, int cost){
		from = t1.id;
		to = t2.id;
		this.cost = cost;
	}
	
	public int compareTo(Edge e){
		return this.cost - e.cost;
	}
}

class Town {
	int x;
	int y;
	HashMap<Town, Integer> map = new HashMap<Town, Integer>();
	
	public Town(int x, int y){
		this.x = x;
		this.y = y;
	}
	
	public void addCost(Town t, int cost){
		map.put(t, cost);
	}
	
	public Edge[] getEdges(){
		Edge[] result = new Edge[map.size()];
		Set<Town> set = map.keySet();
		Iterator<Town> it = set.iterator();
		int i = 0;
		while(it.hasNext()){
			Town tmpTown = it.next();
			result[i] = new Edge(this, tmpTown, map.get(tmpTown));
			i++;
		}
		
		return result;
	}
}