프로그래머스 - 지형 이동

https://programmers.co.kr/learn/courses/30/lessons/62050

코딩테스트 연습 - 지형 이동

BFS와 완전 탐색을 응용해서 풀 수도 있는 문제입니다. 그러나 이 방법은 효율성 테스트를 통과하지 못하고 시간초과를 내기 때문에 최소의 비용으로 모든 Node를 연결하는 방법을 찾는 최소 신장 트리(Minumn Spanning Tree)를 이용해 문제를 풀겠습니다.

 

입출력 예에서도 볼 수 있듯이, 가장 먼저 해야할 것은 연결되어 있는 지형을 하나의 Node로 판별해야 합니다. 저는 번호를 붙여서 판별하겠습니다.

 

  w, h = len(land), len(land[0])
    answer = 0
    visited = [[0] * h for _ in range(w)]
    dx, dy = [0, 1, 0, -1], [1, 0, -1, 0]

    num = 0
    for i in range(w):
        for j in range(h):
            if not visited[i][j]:
                q = [[i, j]]
                visited[i][j] = num + 1
                while q:
                    x, y = q[0][0], q[0][1]
                    for k in range(4):
                        kx, ky = x + dx[k], y + dy[k]
                        if 0 <= kx < h and 0 <= ky < w:
                            if abs(land[x][y] - land[kx][ky]) <= height:
                                if visited[kx][ky] == 0:
                                    visited[kx][ky] = num + 1
                                    q.append([kx, ky])

                    q.pop(0)
                num += 1

맨 처음 0으로 초기화 되어 있는 visited 배열에 번호를 붙입니다. 그 번호는 num 변수에 저장되어 있고, BFS로 높이 차이가 height 이상 나지 않는 칸이라면 전부 같은 번호를 매겨줍니다.

 

번호가 전부 매겨졌다면 같은 번호를 가진 지형을 하나의 Node로 보고, 최소의 비용으로 연결되는 Edge를 찾아 정보를 저장합니다.

 

    weight_dict = {}
    for i in range(w):
        for j in range(h):
            for k in range(4):
                x, y = i + dx[k], j + dy[k]
                if 0 <= x < h and 0 <= y < w:
                    if visited[i][j] != visited[x][y]:
                        l1, l2 = visited[i][j], visited[x][y]
                        weight_dict[l1, l2] = min(weight_dict.get((l1, l2), 10000), abs(land[i][j] - land[x][y]))

 

인접한 지형의 번호가 다를 때(if visited[i][j] != visited[x][y]) 그 번호를 tuple로 저장해 weight_dict의 key로 지정합니다. tuple은 순서가 없으므로 (l1, l2)의 값을 계속 최솟값으로 갱신하며 모든 지형을 탐색합니다.

 

모든 Node에 대한 최소 비용의 Edge 정보가 모두 저장되었습니다. 이제 최소 신장 트리를 구현합니다. 저는 Kruskal 알고리즘을 이용하겠습니다. 최소 신장 트리에 관한 개념은 잘 정리된 블로그가 있으니 참고하면 좋겠습니다.

https://gmlwjd9405.github.io/2018/08/28/algorithm-mst.html

[알고리즘] 최소 신장 트리(MST, Minimum Spanning Tree)란 - Heee's Development Blog

 

오름차순 정렬을 위해 dictionary를 list로 만든 뒤 Edge를 기준으로 정렬합니다.

    bridge_num = num - 1
    weight_list = [(value, (key[0], key[1])) for key, value in weight_dict.items()]
    weight_list.sort()
    parents = [x for x in range(num+1)]
    size = [1 for _ in range(num+1)]

모든 Node를 연결했을 때를 확인하기 위해 Edge의 갯수를 저장합니다. Edge의 갯수는 Node의 갯수 -1입니다.

 

해당 Node의 부모가 어떤 Node인지 저장할 parents 리스트가 필요합니다. 일반적인 Tree의 형태는 부모 -> 자식으로 내려가면서 탐색하는 구조이지만 최소 신장 트리에서는 반대입니다.

 

size 배열은 각 Node가 속한 집합의 크기를 저장할 리스트입니다. Tree를 하나로 합칠 때 사이즈가 작은 쪽을 큰 쪽으로 합쳐 Tree 탐색 시간을 줄이기 위해 집합의 크기를 저장합니다.

 

이제 최소 비용의 Edge부터 연결되어 있는 2개의 Node가 같은 집합에 속해있는지 확인합니다. 두 Node의 최종 부모가 같다면 같은 집합에 속해있고 그렇지 않다면 다른 집합에 속해 있으므로 하나의 집합을 다른 쪽으로 합쳐 같은 부모를 가지도록 만듭니다.

 

def find_parent(node, parents):
    if node == parents[node]:
        return parents[node]
    root = find_parent(parents[node], parents)
    return root

for weight in weight_list:
        if not bridge_num: break
        w, n1, n2 = weight[0], weight[1][0], weight[1][1]
        x, y = find_parent(n1, parents), find_parent(n2, parents)
        if x != y:
            if size[x] > size[y]:
                parents[y] = x
                size[x] += size[y]
            else:
                parents[x] = y
                size[y] += size[x]
            bridge_num -= 1
            answer += w

find_parent 함수는 부모가 자기 자신인 Tree의 최상단 Node를 찾을 때까지 재귀적으로 작동합니다. root 변수에 재귀 함수의 return을 계속 할당하며 부모를 찾습니다.

 

def find_parent(node, parents):
    while node != parents[node]:
        parents[node] = parents[parents[node]]
        node = parents[node]
    return parents[node]

위와 같은 방식으로 부모 Node를 부모 Node의 부모 Node로 갱신하고 현재 node를 부모 Node로 갱신하며 최종 부모를 찾아갈 수도 있겠습니다.

 

만약 서로 다른 부모를 갖고 있다면 n1, n2 이 2개의 Node는 다른 집합에 속해있습니다. Node가 속한 집합의 크기를 비교해 작은 쪽의 부모를 큰 쪽으로 지정합니다. 이 방법으로 탐색 시간을 줄일 수 있습니다.

최소 신장 트리를 구성하는 Edge이므로 정답에 Edge의 가중치를 더해주고, bridge의 갯수를 하나 줄여줍니다. bridge의 갯수가 0이 되면 트리가 전부 구성된 것이기 때문에 탐색을 종료합니다.

 

 

전체 코드

import sys
sys.setrecursionlimit(10**8)

def find_parent(node, parents):
    if node == parents[node]:
        return node
    root = find_parent(parents[node], parents)
    return root

def solution(land, height):
    w, h = len(land), len(land[0])
    answer = 0
    visited = [[0] * h for _ in range(w)]
    dx, dy = [0, 1, 0, -1], [1, 0, -1, 0]

    num = 0
    for i in range(w):
        for j in range(h):
            if not visited[i][j]:
                q = [[i, j]]
                visited[i][j] = num + 1
                while q:
                    x, y = q[0][0], q[0][1]
                    for k in range(4):
                        kx, ky = x + dx[k], y + dy[k]
                        if 0 <= kx < h and 0 <= ky < w:
                            if abs(land[x][y] - land[kx][ky]) <= height:
                                if visited[kx][ky] == 0:
                                    visited[kx][ky] = num + 1
                                    q.append([kx, ky])

                    q.pop(0)
                num += 1

    weight_dict = {}
    for i in range(w):
        for j in range(h):
            for k in range(4):
                x, y = i + dx[k], j + dy[k]
                if 0 <= x < h and 0 <= y < w:
                    if visited[i][j] != visited[x][y]:
                        l1, l2 = visited[i][j], visited[x][y]
                        weight_dict[(l1, l2)] = min(weight_dict.get((l1, l2), 10000), abs(land[i][j] - land[x][y]))

    bridge_num = num - 1
    weight_list = [(value, (key[0], key[1])) for key, value in weight_dict.items()]
    weight_list.sort()
    parents = [x for x in range(num+1)]
    size = [1 for _ in range(num+1)]
    for weight in weight_list:
        if not bridge_num: break
        w, n1, n2 = weight[0], weight[1][0], weight[1][1]
        x, y = find_parent(n1, parents), find_parent(n2, parents)
        if x != y:
            if size[x] > size[y]:
                parents[y] = x
                size[x] += size[y]
            else:
                parents[x] = y
                size[y] += size[x]

            bridge_num -= 1
            answer += w
    return answer

실행 결과

채점을 시작합니다.

정확성 테스트

테스트 1 〉	통과 (0.09ms, 10.7MB)
테스트 2 〉	통과 (0.17ms, 10.8MB)
테스트 3 〉	통과 (0.12ms, 10.8MB)
테스트 4 〉	통과 (0.13ms, 10.8MB)
테스트 5 〉	통과 (0.12ms, 10.8MB)
테스트 6 〉	통과 (0.14ms, 10.8MB)
테스트 7 〉	통과 (0.16ms, 10.8MB)
테스트 8 〉	통과 (0.16ms, 10.8MB)
테스트 9 〉	통과 (0.21ms, 10.8MB)
테스트 10 〉	통과 (0.19ms, 10.8MB)
테스트 11 〉	통과 (0.17ms, 10.8MB)
테스트 12 〉	통과 (0.25ms, 10.8MB)
테스트 13 〉	통과 (0.22ms, 10.7MB)
테스트 14 〉	통과 (0.41ms, 10.9MB)
테스트 15 〉	통과 (2.06ms, 10.8MB)
테스트 16 〉	통과 (21.60ms, 12.5MB)
테스트 17 〉	통과 (93.32ms, 19.2MB)
테스트 18 〉	통과 (86.43ms, 18.5MB)
테스트 19 〉	통과 (82.27ms, 16.5MB)
테스트 20 〉	통과 (463.19ms, 54.4MB)
테스트 21 〉	통과 (468.78ms, 54.5MB)
테스트 22 〉	통과 (1310.07ms, 98.3MB)
테스트 23 〉	통과 (1232.33ms, 101MB)
테스트 24 〉	통과 (272.35ms, 77.4MB)
테스트 25 〉	통과 (288.33ms, 77.6MB)
테스트 26 〉	통과 (1129.22ms, 94.7MB)
테스트 27 〉	통과 (1146.95ms, 95MB)
테스트 28 〉	통과 (1270.18ms, 92.5MB)
테스트 29 〉	통과 (1048.72ms, 90.1MB)
테스트 30 〉	통과 (1053.78ms, 97.8MB)

채점 결과

정확성: 100.0

합계: 100.0 / 100.0