티스토리 뷰
n개의 노드가 있는 그래프가 있습니다. 각 노드는 1부터 n까지 번호가 적혀있습니다. 1번 노드에서 가장 멀리 떨어진 노드의 갯수를 구하려고 합니다. 가장 멀리 떨어진 노드란 최단경로로 이동했을 때 간선의 개수가 가장 많은 노드들을 의미합니다.
노드의 개수 n, 간선에 대한 정보가 담긴 2차원 배열 vertex가 매개변수로 주어질 때, 1번 노드로부터 가장 멀리 떨어진 노드가 몇 개인지를 return 하도록 solution 함수를 작성해주세요.
노드의 개수 n은 2 이상 20,000 이하입니다.
간선은 양방향이며 총 1개 이상 50,000개 이하의 간선이 있습니다.
vertex 배열 각 행 [a, b]는 a번 노드와 b번 노드 사이에 간선이 있다는 의미입니다.
그래프의 1번 노드에서부터 가장 멀리 떨어진 노드의 수를 구하는, 그래프 문제이다.
- 처음 짠 코드
이걸 어떻게 풀어야 하나, 하면서 그래프랑 관련된 알고리즘들을 생각해보다가 제일 먼저 떠오른게 다익스트라 알고리즘이다. 다익스트라 알고리즘은 하나의 정점으로부터 다른 모든 정점까지의 최단 경로를 구하는 (single-source & multi-destination) 알고리즘이기에 마침 이 문제가 요구하는 조건이랑 딱 맞다.
보통 다익스트라 알고리즘이라고 하면 edge마다 거리 (weight) 가 정해져 있을 때를 예시로 많이 들지만, 모든 edge의 weight가 1이라고 생각하고 그냥 쓰면 되지 않을까? 싶었다.
각 정점까지의 최단 경로를 구한 다음 그 중 max값을 구하고, 그 max값과 같은 값을 가진 정점이 몇 개인지 count하면 된다.
...근데 다익스트라 알고리즘 어떻게 구현하더라? 수도 없이 배웠던 거지만 또 까먹어서 위키백과에 있는 의사코드를 참조했다.
짧게 쓰면 dist라는 배열에 최단거리를 저장하는데 처음엔 다 INF (코드에선 대충 적당히 큰 값 999999) 이다. 1번 노드부터 시작해서 이웃 노드들을 방문하면서 dist 배열의 값들을 갱신해 나간다. 기존에 저장되어있던 값보다 새롭게 계산한 값이 더 작으면 갱신한다.
def solution(n, edge):
dist = [999999 for x in range(n+1)];
prev = [-1 for x in range(n+1)];
Q = set(); # 아직 미방문상태인 정점들을 넣는 집합
for i in range(1, n+1):
Q.add(i); # 모든 정점들을 집합 Q에 넣는다
dist[1] = 0; # 1번 노드부터 시작하므로, 1번 노드와의 거리는 0으로 초기화
while len(Q) > 0: # 집합 Q가 빌 때까지...
mindist = 99999;
u = -1;
for i in range(1, n+1): # 미방문 정점들 중 출발점으로부터의 거리가 가장 짧은 정점 찾기
if i not in Q:
continue;
if dist[i] < mindist:
mindist = dist[i];
u = i;
Q.remove(u); # 그 정점을 방문하고 Q에서 지운다
neighbor = []
for e in edge: # 정점 u의 이웃들을 찾는다
if e[0] == u:
neighbor.append(e[1]);
elif e[1] == u:
neighbor.append(e[0]);
for nei in neighbor: # 정점 u의 이웃들을 하나씩 보면서
alt = dist[u] + 1; # 다음 정점까지의 거리 = 지금까지의 거리 + 1
if alt < dist[nei]: # 이게 기존 값보다 더 작으면 갱신
dist[nei] = alt;
prev[nei] = u;
del dist[0]; # 편의상 0번 정점에는 999999를 저장해놨으므로...
max_dist = max(dist);
answer = dist.count(max_dist);
return answer이 코드는... 처음 몇 개 테스트케이스는 통과하다가 뒤쪽에서 시간초과가 났다.
- 두번째 코드: 우선순위 큐 (힙) 를 이용하는 다익스트라 알고리즘
다익스트라 알고리즘에 대해 찾아보다보니 기본적인 다익스트라 알고리즘 (위의 코드) 은 시간복잡도가 n^2이고 우선순위 큐 (Priority Queue) 를 사용하면 더 빠르게 구현할 수 있다는 설명이 많이 있었다.
우선순위 큐는 말 그대로 큐인데 '우선순위' 라는 특정한 값에 따라 자동으로 정렬이 돼서 우선순위가 제일 낮은게 먼저 pop되는 자료구조이다.
이걸 또 직접 구현해야하나... 라는 생각을 잠깐 했지만 그럴 리가 없다. 파이썬에는 미리-구현된 우선순위 큐 (즐겁다! 마참내) 가 존재하기 때문이다.
from queue import PriorityQueue
pq = PriorityQueue();
pq.put((dist[i], i)); // 그냥 숫자만 넣으면 그게 그대로 우선순위가 되고, 이렇게 튜플 형태로 넣으면 앞의 값이 우선순위가 된다
pq.get(); // 우선순위가 제일 낮은 (우선순위, 노드번호) 를 리턴
이 우선순위 큐를 쓰면 어디가 빨라지냐면 while문 시작하는 부분에 있는 '미방문 상태 정점들 중에서 제일 dist값이 작은 정점을 찾는' 부분이다. 위의 코드처럼 그냥 for문을 쓰면 노드 수만큼 도니까 시간복잡도에 N이 추가되지만, 우선순위 큐 (힙 heap이라고도 한다) 에서 삽입/삭제를 하는데 걸리는 시간은 logN이다.
from queue import PriorityQueue
def solution(n, edge):
pq = PriorityQueue();
dist = [99999 for x in range(n+1)];
prev = [-1 for x in range(n+1)];
dist[1] = 0;
for i in range(1, n+1):
pq.put((dist[i], i));
while pq.qsize() > 0:
pq_get = pq.get();
u = pq_get[1];
if pq_get[0] == 99999: # 큐에 INF만 남으면 끝낸다
break;
neighbor = []
for e in edge:
if e[0] == u:
neighbor.append(e[1]);
elif e[1] == u:
neighbor.append(e[0]);
for nei in neighbor:
alt = dist[u] + 1;
if alt < dist[nei]:
dist[nei] = alt;
prev[nei] = u;
pq.put((alt, nei)) # 큐에 갱신된 값 삽입
del dist[0];
max_dist = max(dist);
answer = dist.count(max_dist);
return answer다른 부분은 크게 바뀐게 없고 집합 Q가 우선순위 큐 pq로 바뀌었다. 집합 Q에는 노드 번호만 들어갔었는데 pq에는 노드 번호와 함께 dist의 값 (=우선순위) 이 들어간다. 기존에는 Q 안에 있는 노드들 중에서 dist 값이 제일 작은걸 찾는 for문을 돌려야 했는데 이렇게 바꾸니까 pq.get() 만으로 해결이 된다.
다만 유의해야 했던 점이 하나 있는데, 원래 이 알고리즘의 의사코드를 보면 dist 값을 갱신하는 부분에서 우선순위 큐 안에 들어있는 값의 우선순위를 같이 갱신하는 게 있다. 근데 이 미리-구현된 PriorityQueue에는 그런 기능을 하는 메서드가 없다. 그래서 그냥 새로 집어넣되, 이러면 처음에 넣어놨던 우선순위가 INF인 애들이 그대로 남기 때문에 while문 중간에 pq에서 INF가 튀어나오면 다 끝난 것으로 판단하고 빠져나오게 만들었다.
하지만 이 코드도 통과하지 못했다. 다익스트라를 설명하는 글에서 우선순위 큐보다 더 빠른 알고리즘에 대한 자세한 설명을 보진 못했는데... 첨부터 잘못 생각했나? 다익스트라를 쓰는게 아닌가?
- 세번째 코드: 이웃을 미리 저장해두기
저 코드에서 어떤 부분이 또 쓸데없이 시간을 잡아먹는지 생각을 해봤다. 코드를 몇번이고 보다보니 의사코드에는 한줄로 써있는데 내 코드에선 for문이 들어가있는 부분이 눈에 띄었다.
바로 edge로부터 neighbor를 찾는 부분이다.
생각해보니 이게 또 for문이라서 시간을 잡아먹는다. 심지어 N만큼 시간을 잡아먹는 것도 아니고 edge 수만큼이다... 굳이 저 로직 안에서 neighbor를 매번 구해야하나? edge는 처음부터 주어지니까 그냥 각 노드의 neighbor 목록을 처음부터 만들어두고 들어가도 될 것 같다.
어떤 자료구조를 쓰면 될까? 1번 노드는 2번 노드랑 연결, 2번 노드는 1, 3, 4, 5번이랑 연결... 전에 썼던 포스팅 중에 비슷하게 그래프 문제를 풀기 위해서 인접 리스트라는걸 만들었던 기억이 난다. 그때처럼 dictionary를 써서 각 노드별 neighbor를 저장하면 될 것 같다.
from queue import PriorityQueue
from collections import defaultdict
def solution(n, edge):
pq = PriorityQueue();
dist = [99999 for x in range(n+1)];
prev = [-1 for x in range(n+1)];
dist[1] = 0;
neighbors = defaultdict(lambda: []);
for e in edge: # 노드번호는 key, 이웃한 노드들의 list가 value
neighbors[e[0]].append(e[1]);
neighbors[e[1]].append(e[0]);
for i in range(1, n+1):
pq.put((dist[i], i));
while pq.qsize() > 0:
pq_get = pq.get();
u = pq_get[1];
if pq_get[0] == 99999:
break;
for nei in neighbors[u]:
alt = dist[u] + 1;
if alt < dist[nei]:
dist[nei] = alt;
prev[nei] = u;
pq.put((alt, nei))
del dist[0];
max_dist = max(dist);
answer = dist.count(max_dist);
return answer코드를 보면 while 안에 존재했던 for문이 밖으로 빠져나오면서 한번만 실행하게 되었다. while 안에서는 neighbors[u] 로 간단하게 미리 만들어둔 dict에 접근할 수 있다.
결과는 성공! 오랜만에 다익스트라 알고리즘도 복습하고 PriorityQueue에 대해서도 알게 해준 문제였다. 이전 포스팅에서 썼던 defaultdict와 lambda도 다시 한번 써먹을 수 있었다.