백준 16957 <체스판 위의 공> Python

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

---

그냥 시뮬레이션으로 풀 수도 있지만 조금 더 효율적으로 풀 수 있는 방법이 존재합니다.

생각해보면 모든 위치에 있는 공은 몇 개의 위치로 모일 수 밖에 없습니다. 울퉁불퉁한 콘크리트 바닥에 물을 뿌리면 물이 고이는 위치가 존재하는 것 처럼 다른 위치로 이동할 수 없는 위치로 공이 모두 고이게 됩니다.

 

해당 위치를 찾기 위해 조건에 따라 그래프를 만듭니다. 어차피 모든 노드에서 다른 노드로 이동할 수 있는 경로는 하나뿐입니다. 그러면 우리는 그래프를 사용해 공이 이동하는 경로를 알 수 있습니다.

 

또 B -> A 의 경로가 있고 C -> B 의 경로가 존재한다면 C -> A 의 경로로 병합된다는 점을 고려해 Union-Find 알고리즘을 사용한다면 연결되어 있는 모든 노드를 종착 노드로 병합할 수 있습니다.

물론 기존 Union-Find 알고리즘에서 루트 노드가 작은 쪽으로 병합하는 것을 종착 노드로 병합되게 합니다.

예를 들어 B->A, C->B 의 경로가 존재하면 B 노드는 A 그룹으로 병합되고, C 노드는 B 그룹으로 병합되어 결국 C 노드도 A 그룹으로 병합되게 합니다.

 

해당 방법을 통해 체스판 내부의 모든 공간에서 종착지를 알게 되었으니 개수를 세어 출력만 해주면 됩니다.

---

import sys

input = sys.stdin.readline


## Union-Find
# Union
def Union(node1, node2, group_list):
    # 각 노드의 그룹
    group1 = Find(node1, group_list)
    group2 = Find(node2, group_list)
    # 두 그룹이 같으면
    if group1 == group2:
        # 병합하지 않음
        return False, group_list
    
    # 두 그룹이 다르면
    # 타겟 쪽으로 병합
    group_list[group1] = group2
    
    return True, group_list


# Find
def Find(node, group_list):
    # 그룹의 대표가 자신이 아니면
    if group_list[node] != node:
        # 재귀적으로 업데이트
        group_list[node] = Find(group_list[node], group_list)
    
    return group_list[node]


def solution(R, C, nums):
    # 이동 방향
    move_dirs = [[-1, -1], [-1, 0], [-1, 1], [0, 1], [1, 1], [1, 0], [1, -1], [0, -1]]
    # 그래프 / graph[i][j] = [i, j] 에서 이동할 위치
    graph = [[[] for _ in range(C)] for _ in range(R)]
    for i in range(R):
        for j in range(C):
            # 가장 작은 인덱스
            min_idx = []
            # 가장 작은 값
            min_num = nums[i][j]
            # 이동 방향 순회
            for x, y in move_dirs:
                # 타겟 위치
                ii = i+x
                jj = j+y
                # 인덱스 범위 내부면
                if (ii >= 0 and ii < R) and (jj >= 0 and jj < C):
                    # 가장 작은 값 / 인덱스 갱신
                    if nums[ii][jj] < min_num:
                        min_num = nums[ii][jj]
                        min_idx = [ii, jj]
            
            # 그래프 갱신
            graph[i][j] = min_idx
    
    # 그룹 리스트 
    group_list = list(range(R*C))
    # 그래프 따라서 병합
    for i in range(R):
        for j in range(C):
            # 타겟 인덱스
            # 존재하지 않으면 패스
            if not graph[i][j]:
                continue

            ti, tj = graph[i][j]
            state, group_list = Union(i*C + j, ti*C + tj, group_list)

    # 결과
    result = [[0 for _ in range(C)] for __ in range(R)]
    for g in group_list:
        # 그룹
        group = Find(g, group_list)
        # 인덱스
        i = group // C
        j = group % C
        # 결과 갱신
        result[i][j] += 1
    
    for r in result:
        print(*r)


# 입력
R, C = map(int, input().strip().split())
nums = [list(map(int, input().strip().split())) for _ in range(R)]

solution(R, C, nums)