[プログラマ]配信

問題を理解する

問題を解く

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class Main {

	public static List<int[]>[] g = new List[51]; // graph

	public static int[] table = new int[51]; // 최소 비용 테이블

	public static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
	// public static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
	public static StringTokenizer st;

	public static void setUp() throws IOException {

	}

	// road[i] 는 원소의 개수가 3 이다
	// 두 마을 a,b를 지나는 도로가 여러개라고 하더라도, 어차피 PQ 는 min heap 으로 만들거라.. 그 중 최소비용을 사용하게 될 거임
	public int solution(int N, int[][] road, int K) {
		int answer = 0;

		init(N, road);
		dikstra();

		answer = (int)Arrays.stream(table)
		  .filter(n -> n <= K)
		  .count();

		return answer;
	}

	public void init(int n, int[][] road) {
		for (int i = 1; i <= n; i++) {
			g[i] = new ArrayList<>();
		}

		for (int[] p : road) {
			g[p[0]].add(new int[] {p[1], p[2]});
			g[p[1]].add(new int[] {p[0], p[2]});
		}

		Arrays.fill(table, Integer.MAX_VALUE);

		table[1] = 0;

	}

	public void dikstra() {
		PriorityQueue<int[]> pq = new PriorityQueue<>((o1, o2) -> o1[1] - o2[1]);
		pq.add(new int[] {1, 0});

		while (!pq.isEmpty()) {
			int[] cur = pq.poll(); // 노드이름, 최소비용
			// 현재 최소비용 테이블의 비용보다 크면 패쓰
			if (table[cur[0]] < cur[1])
				continue;

			// 현재 노드의 인접노드 - adj[0] 노드 이름, adj[1] 현재노드 -> 그 노드 비용
			for (int[] adj : g[cur[0]]) {
				int nCost = cur[1] + adj[1]; // 현재 노드를 거쳐 인접노드로 갈 때 비용
				if (nCost >= table[adj[0]])
					continue;
				table[adj[0]] = nCost;
				pq.add(new int[] {adj[0], nCost});
			}
		}
	}

	public static void main(String[] args) throws IOException {
		Main main = new Main();
		int res = main.solution(5, new int[][] {{1, 2, 1}, {2, 3, 3}, {5, 2, 2}, {1, 4, 2}, {5, 3, 1}, {5, 4, 2}},
		  3);
		System.out.println(res);
		res = main.solution(6,
		  new int[][] {{1, 2, 1}, {1, 3, 2}, {2, 3, 2}, {3, 4, 3}, {3, 5, 2}, {3, 5, 3}, {5, 6, 1}},
		  4);
		System.out.println(res);
	}
}