728x90
문제
https://school.programmers.co.kr/learn/courses/30/lessons/12978
코드
자바
import java.util.*;
class Edge implements Comparable<Edge> {
public int vex;
public int cost;
public Edge(int vex, int cost) {
this.vex = vex;
this.cost = cost;
}
@Override
public int compareTo(Edge ob) {
return this.cost - ob.cost;
}
}
class Solution {
public int solution(int N, int[][] road, int K) {
ArrayList<ArrayList<Edge>> graph = new ArrayList<>();
for (int i = 0; i <= N; i++) {
graph.add(new ArrayList<Edge>());
}
for (int[] r : road) {
graph.get(r[0]).add(new Edge(r[1], r[2]));
graph.get(r[1]).add(new Edge(r[0], r[2]));
}
int[] distance = new int[N + 1];
Arrays.fill(distance, Integer.MAX_VALUE);
PriorityQueue<Edge> pQ = new PriorityQueue<>();
pQ.offer(new Edge(1, 0));
distance[1] = 0;
while (!pQ.isEmpty()) {
Edge tmp = pQ.poll();
if (distance[tmp.vex] >= tmp.cost) {
for (Edge ob : graph.get(tmp.vex)) {
if (distance[ob.vex] > tmp.cost + ob.cost) {
distance[ob.vex] = tmp.cost + ob.cost;
pQ.offer(new Edge(ob.vex, distance[ob.vex]));
}
}
}
}
int answer = 0;
for (int d : distance) {
if (d <= K)
answer++;
}
return answer;
}
}
파이썬
import heapq
INF = int(1e9)
def solution(N, road, K):
graph = [[] for _ in range(N + 1)]
distance = [INF] * (N + 1)
for a, b, c in road:
graph[a].append((b, c))
graph[b].append((a, c))
queue = []
heapq.heappush(queue, (0, 1))
distance[1] = 0
while queue:
dist, now = heapq.heappop(queue)
if distance[now] < dist:
continue
for i in graph[now]:
cost = dist + i[1]
if cost < distance[i[0]]:
distance[i[0]] = cost
heapq.heappush(queue, (cost, i[0]))
return len([i for i in distance if i <= K])
728x90
'2022 하계방학 코테 특강 > Graph, Tree 알고리즘' 카테고리의 다른 글
[프로그래머스] 가장 먼 노드 (0) | 2022.09.18 |
---|---|
플로이드 와샬 알고리즘 (0) | 2022.09.18 |
다익스트라 알고리즘 (0) | 2022.09.18 |