[Programmers] 지형 이동(BFS, MST)

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

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

 

from collections import deque, defaultdict
import math
import sys
 
sys.setrecursionlimit(10 ** 6)
 
# Union Find 함수
def find_parent(x, parent):
    if x == parent[x]:
        return x
    else:
        p = find_parent(parent[x], parent)
        parent[x] = p
        return p
 
# Union Find 함수
def union(x, y, parent):
    x = find_parent(x, parent)
    y = find_parent(y, parent)
    parent[y] = x
 
# 처음 그룹을 나누기 위한 BFS
def bfs(land, start, visited, height, group):
    dirs = [(0, 1), (0, -1), (1, 0), (-1, 0)]
    queue = deque()
    queue.append(start)
    while queue:
        y, x = queue.popleft()
        visited[y][x] = group
        for dy, dx in dirs:
            ny, nx = y + dy, x + dx
            if 0 <= ny < len(land) and 0 <= nx < len(land[0]) and visited[ny][nx] == 0 and abs \
                        (land[ny][nx] - land[y][x]) <= height:
                visited[ny][nx] = group
                queue.append((ny, nx))
 
# 인접한 블록이 그룹이 다른 경우 사다리를 놓을 후보지에 추가
def find_height(visited, maps, table):
    dirs = [(0, 1), (0, -1), (1, 0), (-1, 0)]
    for y in range(len(maps)):
        for x in range(len(maps[0])):
            rx = x + 1
            dy = y + 1
 
            if rx < len(maps[0]) and visited[y][x] != visited[y][rx]:
                a, b = min(visited[y][x], visited[y][rx]), max(visited[y][x], visited[y][rx])
                table[(a, b)] = min(table[(a, b)], abs(maps[y][x] - maps[y][rx]))
 
            if dy < len(maps) and visited[dy][x] != visited[y][x]:
                a, b = min(visited[dy][x], visited[y][x]), max(visited[dy][x], visited[y][x])
                table[(a, b)] = min(table[(a, b)], abs(maps[dy][x] - maps[y][x]))
 
 
def solution(land, height):
    visited = [[0 for _ in range(len(land[0]))] for _ in range(len(land))]
    group = 1
 
    for y in range(len(land)):
        for x in range(len(land[0])):
            if visited[y][x] == 0:
                bfs(land, (y, x), visited, height, group)
                group += 1
 
    table = defaultdict(lambda: math.inf)
    find_height(visited, land, table)
    # 사다리 놓을 후보를 작은 순서로 정렬, 앞에서 빼내면서 담으면 답이 된다
    table = sorted(table.items(), key=lambda x: x[1])
 
    answer = 0
    nodes = {i: i for i in range(1, group)}
    for (a, b), value in table:
        # 두 그룹이 같은 그룹인지 판단
        if find_parent(a, nodes) != find_parent(b, nodes):
            union(a, b, nodes)
            answer += value
 
        if len(nodes.values()) == 1:
            return answer
 
    return answer
 
 
land = [[1, 4, 8, 10], [5, 5, 5, 5], [10, 10, 10, 10], [10, 10, 10, 20]]
height = 3
solution(land, height)