[SWEA] D4 : 하나로(3459) Python

이 문제는 인도네시아의 섬들을 모두 연결하는데 환경부담금이 가장 적게 나오도록 연결하는 방법을 묻는 문제이다. 즉 가중치(환경부담금)를 이용한 최소 신장 트리를 구하는 문제이다. 조금 더 쉽게 말하자면 부담금이 가장 적게 나오는 방법으로 모든 섬을 한점 잇기 하는 방법을 찾는 방법이다.

이 문제에서는 X좌표와 Y좌표, 환경부담세율을 Input값으로 각 거리에 비례하는 환경부담금이 최소가 되는 최소 신장 트리를 구하는 문제이다. 문제에 대한 설명은 이 정도로 하고 내가 이문제를 풀기 위해 사용한 Prim알고리즘에 대해 잠깐 알아보고 코드를 정리해 보겠다.

MST(Minimum Spanning Tree)를 구하는 알고리즘 중 대표적으로 Prim과 Kruskal 알고리즘이 있는데, 두개다 자신이 없었던 나는 이 문제를 풀면서 프림 알고리즘을 복습해 보았다.

프림알고리즘 설명 (출처 : 위키트리)

먼저 위키트리에서 설명하는 프림 알고리즘의 설명을 보면 너무어렵다. 알고리즘을 공부할 때 코드를 보면서 교수님의 설명을 들으면 그나마 이해 가능한데(그래도 어렵다), 정의나 위키백과 등을 찾아보면 이해하기가 하늘의 별따기인 것 같다. 

그래서 나도 프림알고리즘을 이해하기 위해 교수님의 코드를 보면서 나만의 이해방법으로 정리하고 있다. 이 문제를 통해서 프림알고리즘을 한번 더 정리해보자!

일단은 프림알고리즘을 위해 내가 사용한 변수들이다.
1. D : list > 각 섬(index)의 위치로 올 때 계산된 가중치(이 문제에서는 환경부담금)
2. visited : list > 각 섬의 방문 기록체크
3. PI : list > 자신의 부모 섬 인덱스(현재 섬으로 오기 전 섬의 인덱스를 기록)
4. X, Y, E : num > 문제에서 제공하는 수 (X좌표, Y좌표, 환경부 담세율)
5. arr : array > 각 섬에서 섬까지의 거리(가중치)

변수 생성과정에서 한 것은 mst를 계산할 때 편의를 위해 각 섬과 섬의 거리를 행렬(arr)로 만들어 놓았다.

변수생성

그다음으로 mst함수를 호출하여 prim알고리즘을 구현했다. 먼저 코드부터 보자면 이렇게 된다

프림알고리즘

1. 우선 첫 값(섬)은 0으로 시작한다(어느 점에서 시작해도 결과와 상관없음) 첫 위치는 아직 다리를 짓지 않았기 때문에, D [u]는 0이 된다. 
2. 프림 알고리즘은 결국 각 노드를 가장 짧은 거리로 한번 방문하기 때문에 N만큼 과정이 반복되면 끝난다.
3. 그다음은 현재 위치에서 내가 갈 수 있는(활성화된 곳이라고 표현하겠음) 곳 중에 방문하지 않고 가중치가 가장 작은 곳을 선택한다. (물론 처음에는 아직 활성화시키지 않았기 때문에 아무 곳도 선택하지 못함)
4. 선택된 곳은 MST노드 중 한 길이 되니 정답처리와 방문 체크를 한다.
5. 그다음은 선택된 곳을 기준으로 또 활성화를 시켜준다. 여기서 활성화를 시켜준다는 것은 아직 방문하지 않은 곳 중에서 내가 갈 수 있는 (이문제에서는 모든 점에서 모든 섬을 갈 수 있다.) 곳들의 D [v]를 현재 위치에서 v까지의 가중치로 크기 비교 후 최신화(활성화) 해준다.
6. 다시 3번으로 돌아와 활성화된 곳 중 가장 작은 곳을 선택하며 이 과정을 반복한다.

전체코드를 보면,

import math

def mst():
    total = 0
    u = 0
    D[u] = 0

    for i in range(N):
        # 가중치 최솟값 찾기
        min = math.inf
        for v in range(N):
            if v != u and visited[v] == 0 and min > D[v]:
                min = D[v]
                u = v

        visited[u] = 1
        total += arr[PI[u]][u]

        for v in range(N):
            if visited[v] == 0:
                if arr[u][v] < D[v]:
                    D[v] = arr[u][v]
                    PI[v] = u

    return math.floor(total * E + 0.5)


T = int(input()) + 1
for tc in range(1,T):
    N = int(input())
    D = [math.inf] * N
    visited = [0] * N
    PI = list(range(N))
    X = list(map(int,input().split()))
    Y = list(map(int,input().split()))
    E = float(input())
    arr = [[0]* N for _ in range(N)]
    for i in range(N):
        for j in range(N):
            arr[i][j] = arr[j][i] = ((X[i]-X[j])**2+(Y[i]-Y[j])**2)  #가중치 기록

    print(f'#{tc} {mst()}')

이렇게 된다.

다익스트라도 이런 식으로 구현 가능하다는 점에서 프림 알고리즘과 유사하다고 하지만 나는 아직 공부하는 입장이고, 다익스트라를 이렇게 구현하지 않아서, 프림 알고리즘을 이해하는데 조금 어려웠다.

문제 자체는 어렵지 않지만, 알고리즘 자체가 어렵기 때문에 이 문제로 복습하는 것이 많이 도움이 됐다.
오늘도 하루 한 개 알고리즘 파이팅!

문제 출처 : https://swexpertacademy.com/main/code/problem/problemDetail.do?problemLevel=4&contestProbId=AV15StKqAQkCFAYD&categoryId=AV15StKqAQkCFAYD&categoryType=CODE&problemTitle=&orderBy=RECOMMEND_COUNT&selectCodeLang=ALL&select-1=4&pageSize=10&pageIndex=2