99클럽 코테 스터디 10일차 TIL: 다익스트라

오늘의 문제: 18352. 특정 거리의 도시 찾기

오늘 문제는 언뜻 보기에 다익스트라 문제이다. 하지만 이 문제가 실버2인 이유가 있지 않을까 하고 바라보니 단순 BFS로 풀어도 문제가 없지 않을까 하는 생각이 먼저 들어서 BFS로 접근해보기로 했다.

아이디어: BFS로 풀 수 있지 않을까

from collections import deque

n,m,k,x=map(int,input().split())
graph=[[] for i in range(n+1)]
visit=[-1 for i in range(n+1)]

for i in range(m):
    a,b=map(int,input().split())
    graph[a].append(b)

def bfs(graph, start):
    deq=deque()
    idx=0
    visit[start]=idx
    deq.append(start)

    while deq:
        node=deq.popleft()
        idx=visit[node]
        for next in graph[node]:
            if visit[next]!=-1: continue
            else:
                visit[next]=idx+1
                deq.append(next)
    sum=0
    for i in range(1,n+1):
        if visit[i]==k: 
            print(i)
            sum+=1
    if sum==0: print(-1)
    
bfs(graph,x)

BFS 같은 기본적인 그래프는 이제 문제없이 짤 수 있다. 하지만 공부를 하기 위해서 다익스트라로도 짜기로 했다.

아이디어: 다익스트라 간선 비용을 1로 하자.

다익스트라를 공부 안한지 너무 오래 돼서 기억이 나지 않지만, 예전 답안을 보니 생각보다 BFS랑 비슷한 내용이 많았다. 약간 다른 점이라면 deque가 아닌 heapq를 이용한다는 것이다.

비용이 더 같거나 적을 때만 다시 heapq에 넣어서 검사하는 형식이다. 생각보다 이렇게 보니 어렵지 않은 로직이었다. 너무 쫄아있던게 아닌가 싶기도 하다. 다만 heapq를 이용하는 법은 좀 배울만 하다.

import heapq, math

n,m,k,x=map(int,input().split())
graph = [[] for _ in range(n+1)]
result = [ math.inf for _ in range(n+1)]

heap = []

# 그래프 채우기
for _ in range(m):
    a,b=map(int,input().split()) #
    graph[a].append((b,1))
    # graph[b].append((a,1))
    # 현재 그래프에 (다음 노드, 비용)을 추가

def dijkstra(start, res):
    heapq.heappush(heap, (0, start))
    res[start] = 0

    while heap:
        dist, node = heapq.heappop(heap)
        # 거리, 현재 노드

        if res[node] < dist: continue
        # 이미 현재 노드의 결과가 값이 더 작은 경우에는 패스

        for next in graph[node]:
            # 그렇다면 지금 노드에서 다른 노드로 탐색
            cost = dist + next[1]
            # 각 다음 노드에 현재까지의 비용 + 다음 노드로의 비용을 계산함

            if cost < res[next[0]]:
                # 계산 결과가 더 작은 값일 경우에는 갱신
                res[next[0]] = cost
                heapq.heappush(heap, (cost, next[0]))

dijkstra(x, result)
c = []
for i in range(1, n+1):
    if result[i]==k: c.append(i)
if len(c)==0: print(-1)
else: print(*sorted(c))

오늘의 회고

티어를 높이고 싶다면 다익스트라처럼 구현이 쉽지만 티어가 높은 문제들을 풀어야 한다. 다익스트라 문제를 작년에 슥 풀고 1년 정도 관심을 끄고 있었는데 다시 맛있는 문제들을 찾아보면 할만 할 듯하다.