[프로그래머스] 배달

문제

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])