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

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

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

 

파이썬 코드

from collections import deque
from copy import deepcopy

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

# 나무 성장과 번식
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 drug[nx][ny] == 0:
                if graph[nx][ny] > 0:
                    graph[x][y] += 1
                if graph[nx][ny] == 0:
                    count += 1
        temp.append((x, y, count))

    copy_graph = deepcopy(graph)
    while temp:
        x, y, value = temp.pop()
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            if 0 <= nx < n and 0 <= ny < n and drug[nx][ny] == 0 and copy_graph[nx][ny] == 0:
                graph[nx][ny] += copy_graph[x][y] // value

# 최대로 박멸하는 칸 찾기
def maximum_drug(tree):
    Q = deque(tree)
    dx = [1, 1, -1, -1]
    dy = [1, -1, 1, -1]
    temp = []
    while Q:
        x, y = Q.popleft()
        count = 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:
                        count += graph[nx][ny]
                    if graph[nx][ny] <= 0:
                        break
        temp.append((x, y, count))
    return sorted(temp, key=lambda x:(x[2], -x[0], -x[1]))

# 제초제 뿌리기
def spread_drug(x, y):
    Q = deque()
    Q.append((x, y))
    dx = [1, 1, -1, -1]
    dy = [1, -1, 1, -1]
    while Q:
        x, y = Q.popleft()
        drug[x][y] = c
        graph[x][y] = 0
        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:
                    drug[nx][ny] = c
                    if graph[nx][ny] <= 0:
                        break
                    graph[nx][ny] = 0

# 나무 위치 찾기
def cal_tree():
    tree = []
    for i in range(n):
        for j in range(n):
            if graph[i][j] > 0:
                tree.append((i, j))
    return tree

# 제초제 1 제거
def remove_drug():
    for i in range(n):
        for j in range(n):
            if drug[i][j] > 0:
                drug[i][j] -= 1

answer = 0
for T in range(m):
    ftree = cal_tree()
    if len(ftree) == 0:
        break 
    grow(ftree)
    stree = cal_tree()
    maximum = maximum_drug(stree)
    remove_drug()
    tx, ty, result = maximum.pop()
    answer += result
    spread_drug(tx, ty)

print(answer)