[프로그래머스/Python] 지형 이동

1. 문제 / 난이도 : Level 4

Summer/Winter Coding(2019)

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

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

---

2. 풀이 및 코드 설명

처음에는 대충 원하는 대로 구현했다가 참패... 20%빼고 런타임에러를 맞는 기적을 경험했다. 유력한 가설은 아마 시간초과나 메모리 초과가 아닐까...

그 다음 어떻게 짜야할 지 감이 잡히지 않아서 결국 검색의 도움을 빌렸다. 다른 분들의 풀이를 참고해서 아래와 같이 해결하였다.

 

1. 같은 그룹으로 묶일 수 있는 칸끼리 묶기 (BFS)

2. 각 그룹간의 이동 거리 구하기

3. 최소 이동 거리 구하기 (MST)

 

(1) 초기화 및 설정

N = len(land) 
check = [[0 for _ in range(N)] for _ in range(N)] # 각 칸의 그룹

(2) BFS

group = 1
    for i in range(N):
        for j in range(N):
            if not check[i][j]: # 아직 그룹이 정해지지 않은 칸이라면
                BFS(land, height, check, [i, j], group, N)
                group += 1

• N*N의 칸을 모두 돌면서 아직 그룹이 정해지지 않은 칸이 있다면 해당 칸을 기점으로 BFS를 수행

def BFS(land, height, check, loc, group, N):
    direction = [(0, -1), (0, 1), (1, 0), (-1, 0)]
    q = deque()
    q.append(loc)
    check[loc[0]][loc[1]] = group

    while q:
        x, y = q.popleft()
        for dx, dy in direction:
            nx, ny = x + dx, y + dy
            if nx < 0 or nx >= N or ny < 0 or ny >= N or check[nx][ny]: continue
            if abs(land[nx][ny] - land[x][y]) <= height: # 이동하려는 칸과의 높이 차가 작다면
                q.append([nx, ny]) # 큐에 넣기
                check[nx][ny] = group # 그룹 표시

(3) 각 그룹간의 이동거리 계산

def find_ladder(check, land, N):
    direction = [(0, -1), (0, 1), (1, 0), (-1, 0)]
    ladders = defaultdict(lambda: math.inf) # 최대값으로 초기화
    
    for i in range(N):
        for j in range(N):
            current = check[i][j]
            for dx, dy in direction:
                nx, ny = i + dx, j + dy
                if nx < 0 or nx >= N or ny < 0 or ny >= N or check[nx][ny] == current: continue # 같은 그룹에 속해 있다면 건너뜀
                dist = abs(land[i][j] - land[nx][ny])
                ladders[(current, check[nx][ny])] = min(dist, ladders[(current, check[nx][ny])]) # 최소거리로 갱신
    return ladders

• ladders : {(x, y): value} 형태로 되어있는 딕셔너리. x그룹에서 y그룹으로 이동할 때 필요한 사다리의 비용을 가지고 있음

 

(4) MST

def kruskal(ladders, group):
    sum = 0
    roots = {_:_ for _ in range(1, group)} # 정점들 root 표시
    for (x, y), value in ladders:
       if find_root(x, roots) != find_root(y, roots): # root가 같지 않을 경우
           union_find(x, y, roots) # root 연결
           sum += value # 값 갱신
       if len(roots.items()) == 1: return sum # 노드들이 전부 연결된 경우 리턴
    return sum

• 크루스칼 알고리즘으로 MST 구현
• Union-find를 적용해서 각 노드의 루트로 올라가 루트가 같지 않을 경우에 값을 갱신해주었다!

---

3. 전체 코드

from collections import deque, defaultdict
import math

def find_root(x, root):
    if x == root[x]: return x
    else:
        r = find_root(root[x], root)
        root[x] = r
        return r

def union_find(x, y, root):
    x_root = find_root(x, root)
    y_root = find_root(y, root)
    root[y_root] = x_root

def BFS(land, height, check, loc, group, N):
    direction = [(0, -1), (0, 1), (1, 0), (-1, 0)]
    q = deque()
    q.append(loc)
    check[loc[0]][loc[1]] = group

    while q:
        x, y = q.popleft()
        for dx, dy in direction:
            nx, ny = x + dx, y + dy
            if nx < 0 or nx >= N or ny < 0 or ny >= N or check[nx][ny]: continue
            if abs(land[nx][ny] - land[x][y]) <= height: 
                q.append([nx, ny])
                check[nx][ny] = group

def find_ladder(check, land, N):
    direction = [(0, -1), (0, 1), (1, 0), (-1, 0)]
    ladders = defaultdict(lambda: math.inf)
    
    for i in range(N):
        for j in range(N):
            current = check[i][j]
            for dx, dy in direction:
                nx, ny = i + dx, j + dy
                if nx < 0 or nx >= N or ny < 0 or ny >= N or check[nx][ny] == current: continue
                dist = abs(land[i][j] - land[nx][ny])
                ladders[(current, check[nx][ny])] = min(dist, ladders[(current, check[nx][ny])])
    return ladders

def kruskal(ladders, group):
    sum = 0
    roots = {_:_ for _ in range(1, group)}
    for (x, y), value in ladders:
       if find_root(x, roots) != find_root(y, roots):
           union_find(x, y, roots)
           sum += value
       if len(roots.items()) == 1: return sum
    return sum

def solution(land, height):
    answer = 0
    N = len(land)
    check = [[0 for _ in range(N)] for _ in range(N)]
    
    group = 1
    for i in range(N):
        for j in range(N):
            if not check[i][j]: 
                BFS(land, height, check, [i, j], group, N)
                group += 1
    
    ladders = find_ladder(check, land, N)
    ladders = sorted(ladders.items(),key=lambda x:x[1])
    
    answer = kruskal(ladders, group)
    return answer