[코드트리] - 나무 박멸 파이썬 풀이

출처 : https://www.codetree.ai/training-field/frequent-problems/problems/tree-kill-all/description?page=1&pageSize=20

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

파이썬 코드

from collections import deque

n, m, k, c = map(int, input().split())
graph = [list(map(int, input().split())) for _ in range(n)]

#dead : 제초제
dead = [[0]*n for _ in range(n)]

# 1. 인접한 4개의 칸 중 나무가 있는 칸 수만큼 성장
def grow(tree):
    Q = deque(tree)
    dx = [1, -1, 0, 0]
    dy = [0, 0, 1, -1]
    temp = []
    while Q:
        x, y = Q.popleft()
        count = 0
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            if 0 <= nx < n and 0 <= ny < n and graph[nx][ny] > 0:
                count += 1
        temp.append((x, y, count))

    while temp:
        x, y, count = temp.pop()
        graph[x][y] += count

# 2. 번식
def bunsik(tree):
    Q = deque(tree)
    dx = [1, -1, 0, 0]
    dy = [0, 0, 1, -1]
    temp = []
    tmp = []
    while Q:
        x, y = Q.popleft()
        count = 0
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            if 0 <= nx < n and 0 <= ny < n and graph[nx][ny] == 0 and dead[nx][ny] == 0:
                count += 1
                tmp.append((nx, ny))
        for _ in range(len(tmp)):
            nx, ny = tmp.pop()
            temp.append((x, y, nx, ny, count))

    while temp:
        x, y, nx, ny, count = temp.pop()
        graph[nx][ny] += graph[x][y] // count

# 3. 가장 많이 박멸되는 칸 찾기
def find_dead(tree):
    Q = deque(tree)
    dx = [1, 1, -1, -1]
    dy = [1, -1, 1, -1]
    visit = [[0]*n for _ in range(n)]
    while Q:
        x, y = Q.popleft()
        visit[x][y] = graph[x][y]
        for i in range(4):
            for j in range(1, k+1):
                nx = x + (dx[i])*j
                ny = y + (dy[i])*j
                if 0 <= nx < n and 0 <= ny < n:
                    if graph[nx][ny] > 0:
                        visit[x][y] += graph[nx][ny]
                    if graph[nx][ny] == 0 or graph[nx][ny] == -1:
                        break

    return visit

# 4. 해당 칸으로부터 제초제를 뿌리고 해당 칸 나무 제거
def spread_dead(x, y, c):
    Q = deque()
    Q.append((x, y))
    dead[x][y] = c
    graph[x][y] = 0
    dx = [1, 1, -1, -1]
    dy = [1, -1, 1, -1]
    while Q:
        x, y = Q.popleft()
        for i in range(4):
            for j in range(1, k+1):
                nx = x + (dx[i])*j
                ny = y + (dy[i])*j
                if 0 <= nx < n and 0 <= ny < n:
                    dead[nx][ny] = c
                    if graph[nx][ny] == 0 or graph[nx][ny] == -1:
                        break
                    graph[nx][ny] = 0

# 해가 한번 지날 때마다 제초제 1씩 줄이기
def remove_dead():
    for i in range(n):
        for j in range(n):
            if dead[i][j] > 0:
                dead[i][j] -= 1



result = 0
for T in range(1, m+1):
    # 나무 위치 찾기
    tree = []
    for i in range(n):
        for j in range(n):
            if graph[i][j] > 0:
                tree.append((i, j))
    # 1번 성장
    grow(tree)

    tree2 = []
    for i in range(n):
        for j in range(n):
            if graph[i][j] > 0:
                tree2.append((i, j))
    
    # 2번 번식시키기
    bunsik(tree2)

    tree3 = []
    for i in range(n):
        for j in range(n):
            if graph[i][j] > 0:
                tree3.append((i, j))
    
    # 가장 많이 박멸되는 칸 찾기
    V = find_dead(tree3)

    # sx, sy : 제초제를 처음 뿌릴 위치
    sx, sy = 0, 0
    # 제거되는 나무 수
    dead_max = 0

    # for 문으로 차례대로 찾으면 같은 수가 있어도 행과 열이 가장 작은 위치가 저장됨
    for i in range(n):
        for j in range(n):
            if V[i][j] > dead_max:
                sx, sy = i, j
                dead_max = V[i][j]

    # 제거된 나무 수를 result 에 더함
    result += dead_max

    # 한 해가 지났으므로 제초제 1씩 줄이기
    remove_dead()

    # 제초제 뿌리기
    spread_dead(sx, sy, c)


print(result)