[골드 3] 백준 16957 - 체스판 위의 공 (파이썬)

분류 : DFS

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

풀이

각 공이 어디로 떨어지는가가 아니라 각 골짜기로 얼마나 많은 공이 모이는가를 역방향으로 계산한다.

각 칸에서 인접한 8방향 중 가장 낮은 칸으로 공이 흘러간다.

그 흐름을 역방향 그래프로 만들어 두고, 실제로 공이 도달할 수 있는 최종 지점부터 시작해서, 역으로 DFS를 돌며 얼마나 많은 공이 이 칸으로 오는지 세는 방식이다.

1. 입력 처리

R, C = map(int, input().split())
B = [list(map(int, input().split())) for _ in range(R)]

dr = [0, 0, 1, 1, 1, -1, -1, -1]
dc = [1, -1, -1, 0, 1, -1, 0, 1]

B[r][c]는 r행 c열의 높이값을 나타내며, dr, dc는 8방향 이동을 위한 델타값이다.

1. 흐름을 반대로 저장할 adj를 선언한다.

adj = {(r, c) : [] for r in range(R) for c in range(C)}
valley = []

adj[(r, c)]는 이 칸으로 공이 도달하는 출발 위치 리스트를 나타내며, valley의 경우 더 낮은 인접같이 없어 공이 머무르게 되는 칸을 나타낸다.

1. 흐름 경로 설정

for r in range(R):
    for c in range(C):
        min_val = B[r][c]
        who = (-1, -1)
        for i in range(8):
            nr, nc = r + dr[i], c + dc[i]
            if 0 <= nr < R and 0 <= nc < C and B[nr][nc] < min_val:
                min_val = B[nr][nc]
                who = (nr, nc)
        
        if who != (-1, -1):
            adj[who].append((r, c))
        else:
            valley.append((r, c))

각 칸 (r, c)에서 가장 낮은 인접 칸 (nr, nc)를 찾고, 더 낮은 칸이 있으면 (nr, nc)로 공이 흐르므로 adj[(nr, nc)]에 (r, c)를 추가한다.

만약 없는 경우 골짜기로 간주한다.

def dfs(r, c):
    global count
    count += 1
    for (nr, nc) in adj[(r, c)]:
        dfs(nr, nc)

역방향 DFS로써, 현재 valley에서 이 칸으로 유입되는 모든 칸들을 방문하며, count에 누적한다.

for vr, vc in valley:
    count = 0
    dfs(vr, vc)
    answer[vr][vc] = count

각 골짜기마다 DFS를 실행해 얼마나 많은 공이 이 골짜기로 모이는지 카운트 하고, 그 수를 answer[vr][vc]에 기록한다.

그런데 이 풀이의 문제점은 중복 방문이 있다는 점이다.

예를 들어 (2, 2)가 여러 valley에 도달하는 경로 중 중복으로 포함되면, DFS에서 여러번 방문하게 되고, 결과적으로 전체 공 개수가 맞지 않게 된다.

코드

# 백준 16957 - 체스판 위의 공
# 분류 : DFS

from collections import deque

R, C = map(int, input().split())
B = [list(map(int, input().split())) for _ in range(R)]

adj = {(r, c) : [] for r in range(R) for c in range(C)}

dr = [0, 0, 1, 1, 1, -1, -1, -1]
dc = [1, -1, -1, 0, 1, -1, 0, 1]

valley = []

for r in range(R) :
    for c in range(C) :
        min_val = B[r][c]
        who = (-1, -1)
        for i in range(8) :
            nr, nc = r + dr[i], c + dc[i]
            if nr < 0 or R <= nr or nc < 0 or C <= nc :
                continue

            if min_val > B[nr][nc] :
                min_val = B[nr][nc]
                who = (nr, nc)
        
        if who != (-1, -1) :
            adj[who].append((r, c))
        else :
            valley.append((r, c))

count = 0

def dfs(r, c) :
    global count
    count += 1

    for (nr, nc) in adj[(r, c)] :
        dfs(nr, nc)
    
answer = [[0] * C for _ in range(R)]

for vr, vc in valley :
    count = 0
    dfs(vr, vc)
    answer[vr][vc] = count
    
for i in range(R) :
    print(*answer[i])