문제

https://www.acmicpc.net/problem/14938

문제 설명

해당 문제는 한 노드에서 m의 비용으로 갈 수 있는 다른 노드들의 합의 최대를 구하는 문제이다. 즉, 모든 노드들에서 각각 m으로 갈 수 있는 노드들의 합을 구한 다음 그 중에서 최대값을 찾아야한다.

풀이 1 (플로이드 워셜)

모든 노드들에서 다른 노드들로의 최단 거리를 구해야하고 노드 개수가 최대 100개이므로 시간 복잡도가 $O(V^3)$인 플로이드 워셜로 충분히 풀이가 가능해보인다.

items = list(map(int, input().split()))
dist = [[INF] * n for _ in range(n)]

위와 같이 노드들을 입력받고 2차원 dist 배열을 모두 INF로 초기화 해준다.

for _ in range(r):
    a, b, l = map(int, input().split())
    dist[a-1][b-1] = l
    dist[b-1][a-1] = l
    
for i in range(n):
    dist[i][i] = 0

간선 정보를 입력 받는다. 양방향 통행이 가능하므로 서로의 노드로 가는 거리를 모두 l로 변경해준다. 그리고 자기 자신으로의 거리 역시 0으로 변경해준다.

for k in range(n):
    for i in range(n):
        for j in range(n):
            dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j])

플로이드 워셜 알고리즘대로 모든 노드를 중간 지점으로 삼아가며 최단 거리를 계산한다.

answer = 0
for i in range(n):
    result = 0
    for j in range(n):
        if dist[i][j] <= m:
            result += items[j]
    answer = max(answer, result)

이후 각 노드들마다 거리가 m 이하인 다른 노드들의 값을 모두 더한 다음, 그 값들 중 최대값을 찾아주면 정답을 찾을 수 있다!

전체 코드 및 소요 시간

import sys
input = sys.stdin.readline
INF = int(1e9)

n, m, r = map(int, input().split())
items = list(map(int, input().split()))
dist = [[INF] * n for _ in range(n)]

for _ in range(r):
    a, b, l = map(int, input().split())
    dist[a-1][b-1] = l
    dist[b-1][a-1] = l

for i in range(n):
    dist[i][i] = 0

for k in range(n):
    for i in range(n):
        for j in range(n):
            dist[i][j] = min(dist[i][j], dist[i][k] + dist[k][j])

answer = 0
for i in range(n):
    result = 0
    for j in range(n):
        if dist[i][j] <= m:
            result += items[j]
    answer = max(answer, result)

print(answer)

풀이 2 (다익스트라)

해당 문제는 최대 노드 개수가 100개이므로 플로이드 워셜로도 안정적인 풀이가 가능하지만 간선의 개수 r이 최대 100이라는 점을 이용하여 더욱 짧은 소요 시간으로 풀이가 가능하다.

이 문제와 같이 양방향 통행이 가능한 무방향 그래프의 경우에는 간선의 개수가 최대 $V(V-1)/2$개 까지 존재 할 수 있다. 즉, 노드의 개수가 100개라면 최대 4950개의 간선이 존재 할 수 있다.

그렇다면 만약 간선이 최대 개수로 존재한다고 가정을 하고 다익스트라로 해당 문제를 풀면 어떻게 될까? 다익스트라의 시간복잡도는 $O(ElogV)$이고 모든 노드들에 대해서 다익스트라 알고리즘을 수행해야 하므로 최종 시간복잡도는 $O(VElogV)$가 될 것이다. 이 때 $V$는 100, $E$는 4950, $log_2{100}$은 약 6.6 이므로 약 3,267,000번의 연산이 필요하다. 노드가 100개일 때 플로이드 워셜로 풀 경우 $O(V^3)$ 이므로 1,000,000번의 연산이 필요한 것과 비교해보면 대략 3.3배의 연산량을 가지는 것을 볼 수 있다. 이것이 바로 일반적인 경우에는 모든 노드들에서 다른 노드들에 대한 최단거리를 모두 구해야 할 때 노드의 개수가 많지 않다면 플로이드 워셜을 사용해야 하는 이유이다. 플로이드 워셜은 최악의 경우에도 $O(V^3)$을 보장하는 반면, 다익스트라는 간선의 개수에 따라 시간복잡도가 유동적이기 때문이다.

하지만 해당 문제는 무방향 그래프임에도 불구하고 간선의 개수가 최대 100개로 제한 되어 있다. 일반적인 경우가 아닌 것이다. 위 계산에서 $E$를 100으로 바꾸면 66,000으로 연산량이 확 줄어든다. 즉, 해당 문제는 모든 노드들에 대해서 다른 노드들에 대한 최단거리를 구해야 함에도 불구하고 플로이드 워셜보다 다익스트라를 이용하는게 더 효율적인 문제이다.

그럼 해당 문제를 다익스트라로 다시 풀어보자!

distance = [INF] * n
distance[start] = 0
queue = []
queue.append((start, distance[start]))

distance 배열을 INF로 초기화를 해준 뒤 시작 노드를 heapq에 넣어준다.

while queue:
    current, dist = heapq.heappop(queue)
    if distance[current] < dist:
        continue
    for next_, next_dist in graph[current]:
        cost = dist + next_dist
        if distance[next_] > cost:
            distance[next_] = cost
            heapq.heappush(queue, (next_, cost))

heapq에서 최단 거리 노드를 하나씩 꺼낸 후 방문하지 않은 노드들 중 그 노드로 가는 거리가 distance에 기록된 거리보다 짧다면 거리를 갱신하고 heapq에 넣는다.

result = 0
for i in range(n):
    if distance[i] <= m:
        result += items[i]

그 후 거리가 m이하인 노드들의 값을 모두 더한다. 해당 로직을 모든 노드들에서 반복하면 된다.

전체 코드 및 소요 시간

import sys
import heapq
input = sys.stdin.readline
INF = int(1e9)

def dijkstra(start):
    global n
    global m

    distance = [INF] * n
    distance[start] = 0
    queue = []
    queue.append((start, distance[start]))

    while queue:
        current, dist = heapq.heappop(queue)
        if distance[current] < dist:
            continue
        for next_, next_dist in graph[current]:
            cost = dist + next_dist
            if distance[next_] > cost:
                distance[next_] = cost
                heapq.heappush(queue, (next_, cost))

    result = 0
    for i in range(n):
        if distance[i] <= m:
            result += items[i]

    return result

n, m, r = map(int, input().split())
items = list(map(int, input().split()))
graph = [[] for _ in range(n)]

for _ in range(r):
    a, b, l = map(int, input().split())
    graph[a-1].append((b-1, l))
    graph[b-1].append((a-1, l))

answer = 0
for i in range(n):
    answer = max(answer, dijkstra(i))

print(answer)

소요 시간이 360ms에서 48ms에서 엄청 단축된 것을 확인 할 수 있다!

후기

해당 문제는 아무런 응용 없이 기본적인 플로이드 워셜, 다익스트라를 사용하면 돼서 난이도 자체는 높지 않았던 문제이다. 하지만 이 문제는 단순히 어떤 알고리즘으로 풀 수 있냐 없냐를 떠나서 데이터의 제한 사항을 잘 파악하고 가장 빠른 알고리즘은 무엇일까를 고민해야 한다고 느끼게 해준 문제였다.

이전까지는 모든 노드들에서 최단 거리를 구해야 했다면 무지성으로 플로이드 워셜을 갈기고 보았는데 이제는 다익스트라 중 뭐가 더 효율적일지 고민을 해봐야겠다!