백준 16724: 피리 부는 사나이(python)

문제

---

풀이

 

지도의 각 칸을 노드로, 성우가 설정한 방향을 간선으로 생각하고 union-find를 사용해 유향 그래프 상 집합을 판별하여 풀이하는 문제다. 각 집합 당 하나의 SAFE ZONE을 배치한 것이 회원들이 지도 어느 구역에 있더라도 SAFE ZONE에 들어갈 수 있는 최소 개수이다.

N, M = map(int, input().split())
maps = []
for _ in range(N):
    maps.append(list(map(str, input().strip())))

입력은 maps 배열에 저장해 주었다. 이후 저장된 UDLR을 간선으로 변환한다.

idx = 0
graph = [[] for _ in range(N * M)]

for i in range(N):
    for j in range(M):
        if maps[i][j] == 'U':
            graph[idx].append(idx - M)	# 1행 아래로 이동
        elif maps[i][j] == 'D':
            graph[idx].append(idx + M)	# 1행 위로 이동
        elif maps[i][j] == 'L':
            graph[idx].append(idx - 1)	# 1열 왼쪽으로 이동
        else:
            graph[idx].append(idx + 1)	# 1열 오른쪽으로 이동
        idx += 1

위와 같이 생각할 수 있는 셈이다

 

이후는 일반적인 유니온 파인드 문제와 유사하다.

parent = [i for i in range(N * M)]

def find(v):
    if parent[v] != v:
        parent[v] = find(parent[v])
    return parent[v]

def union(u, v):
    ru, rv = find(u), find(v)

    if ru != rv:
        parent[ru] = rv

union(), find() 함수를 구현하고

for i in range(N * M):
    for j in graph[i]:
        union(i, j)

for i in range(N * M):  # 모든 원소에 대해 최종 부모 탐색
    parent[i] = find(i)

for문을 순회하며 각 노드에 연결된 노드에 대해 union() 연산을 수행한다. 

 

위 과정을 거치면 parent 배열에는 각 노드의 루트 노드들이 저장되게 되며, parent 배열의 원소 개수가 그래프 내 집합의 수, 즉 필요한 최소 SAFE ZONE의 개수가 된다. parent 원소의 개수를 세기 위해 set()을 사용하였다.

print(len(set(parent)))

 

전체 코드

import sys
input = sys.stdin.readline

N, M = map(int, input().split())
maps = []
for _ in range(N):
    maps.append(list(map(str, input().strip())))

idx = 0
graph = [[] for _ in range(N * M)]

for i in range(N):
    for j in range(M):
        if maps[i][j] == 'U':
            graph[idx].append(idx - M)
        elif maps[i][j] == 'D':
            graph[idx].append(idx + M)
        elif maps[i][j] == 'L':
            graph[idx].append(idx - 1)
        else:
            graph[idx].append(idx + 1)
        idx += 1

parent = [i for i in range(N * M)]

def find(v):
    if parent[v] != v:
        parent[v] = find(parent[v])
    return parent[v]

def union(u, v):
    ru, rv = find(u), find(v)

    if ru != rv:
        parent[ru] = rv

for i in range(N * M):
    for j in graph[i]:
        union(i, j)

for i in range(N * M):  # 모든 원소에 대해 최종 부모 탐색
    parent[i] = find(i)

print(len(set(parent)))