[python] 코드트리 - 나무박멸

문제

https://www.codetree.ai/training-field/frequent-problems/problems/tree-kill-all/description?page=4&pageSize=5

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

 

 

나무의 성장

• 인접 한 네 개의 칸 (상하좌우) 중 나무가 있는 칸의 수만큼 나무 성장
• 모든 나무는 동시에 성장

나무의 번식

• 인접한 4개의 칸 중 벽, 다른 나무, 제초제 모두 없는 칸에 번식 진행
• 각 칸의 나무 그루 수에서 총 번식이 가능한 칸의 개수만큼 나누어진 그루 수만큼 번식되고 나머지는 버림
• 빨간색을 표시된 칸은 파란색으로 표시된 칸에서 번식된 나무임(31//3 = 10)

 

나무의 박멸

• 각 칸 중 제초제를 뿌렸을 때 나무가 가장 많이 박멸되는 칸에 제초제를 뿌림
• 나무가 있는 칸에 제초제를 뿌리게 되면 4개의 대각선 방향으로 k칸만큼 전파
• 전파되는 도중 벽이 있거나 나무가 아얘 없는 칸이 있는 경우, 그 칸 까지는 제초제가 뿌려지며 그 이후의 칸으로는 제초제가 전파되지 않음
• 제초제가 뿌려진 칸에는 c년만큼 제초제가 남아있다가 c+1년째가 될 때 사라짐.
• 제초제가 뿌려진 곳에 다시 제초제가 뿌려지는 경우에는 새로 뿌려진 해로부터 다시 c년동안 제초제가 유지
• k가 2일때 각각의 칸에 제초제가 뿌려진 경우 박멸되는 나무의 총 그루수는 다음과 같음

 

3행 4열이 가장 많은 나무를 박멸시키는 것을 알 수 있으며, 3행 4열에 제초제를 뿌린 이후에는 다음과 같이 변하게 됩니다.

 

 

풀이

문제에서 요구하는대로 순서대로 코드를 작성하면 어렵지 않게 풀수있다. 

1. 나무의 성장 

def grow_tree(arr: List[List[int]]) -> List[List[int]]:
    # ... (코드 내용)

각 칸을 검사하여 나무가 있으면, 주변 4방향의 나무 수를 세고, 그만큼 현재 칸의 나무를 증가시킵니다.

 

2. 나무의 번식

def reproduce_tree(arr: List[List[int]]) -> List[List[int]]:
    # ... (코드 내용)

번식 가능한 칸을 미리 계산하고, 각 나무에 대해 번식할 수 있는 칸의 수를 세어 번식을 처리합니다. 

 

3. 제초제 뿌리기

def exterminate_tree(arr: List[List[int]], k: int, c: int) -> Tuple[List[List[int]], int]:
    # ... (코드 내용)

모든 칸에 대해 제초제를 뿌렸을 때의 효과를 계산하고, 가장 많은 나무를 제거할 수 있는 위치를 찾아 실제로 제초제를 뿌립니다.

대각선 방향으로 k칸만큼 확장하며, 벽이나 격자 경계를 만나면 멈춥니다.

 

4. 제초제 효과 감소

def remove_poison(arr: List[List[int]]) -> List[List[int]]:
    # ... (코드 내용)

모든 칸을 순회하며 제초제가 있는 칸(음수 값)의 유효 기간을 1씩 줄입니다.

 

 

코드

import sys
from typing import List, Tuple

# 방향 상수 (상, 하, 좌, 우)
DX = [0, 0, -1, 1]
DY = [-1, 1, 0, 0]

# 대각선 방향 상수
DIAGONAL_DX = [1, -1, 1, -1]
DIAGONAL_DY = [1, 1, -1, -1]

def is_valid(x: int, y: int, n: int) -> bool:
    """주어진 좌표가 유효한지 확인합니다."""
    return 0 <= x < n and 0 <= y < n

def read_input() -> Tuple[int, int, int, int, List[List[int]]]:
    """입력을 읽어 시뮬레이션 매개변수와 초기 배열을 반환합니다."""
    n, m, k, c = map(int, sys.stdin.readline().split())
    arr = [list(map(int, input().split())) for _ in range(n)]
    return n, m, k, c, arr

def grow_tree(arr: List[List[int]]) -> List[List[int]]:
    """나무를 성장시킵니다."""
    n = len(arr)
    for i in range(n):
        for j in range(n):
            if arr[i][j] <= 0:
                continue
            grow_cnt = 0
            for d in range(4):
                ny = i + DY[d]
                nx = j + DX[d]
                if is_valid(nx, ny, n) and arr[ny][nx] > 0:
                    grow_cnt += 1
            arr[i][j] += grow_cnt
    return arr

def reproduce_tree(arr: List[List[int]]) -> List[List[int]]:
    """나무를 번식시킵니다."""
    n = len(arr)
    visited = [[0] * n for _ in range(n)]
    
    # 나무가 성장할 수 있는 칸 미리 계산
    for i in range(n):
        for j in range(n):
            if arr[i][j] == 0:
                visited[i][j] = 1

    for i in range(n):
        for j in range(n):
            if arr[i][j] < 0 or visited[i][j] == 1:
                continue
            reproduce_list = []
            for d in range(4):
                ny = i + DY[d]
                nx = j + DX[d]
                if is_valid(nx, ny, n) and visited[ny][nx] == 1:
                    reproduce_list.append((ny, nx))
            
            if reproduce_list:
                reproduce_num = arr[i][j] // len(reproduce_list)
                for y, x in reproduce_list:
                    arr[y][x] += reproduce_num
    return arr

def check_remove_tree(arr: List[List[int]], y: int, x: int, k: int) -> int:
    """제거될 나무의 수를 계산합니다."""
    remove_tree_num = arr[y][x]
    n = len(arr)

    for d in range(4):
        for mul in range(1, k + 1):
            nx = x + DIAGONAL_DX[d] * mul
            ny = y + DIAGONAL_DY[d] * mul
            if not is_valid(nx, ny, n) or arr[ny][nx] <= 0:
                break
            remove_tree_num += arr[ny][nx]
    return remove_tree_num

def exterminate_tree(arr: List[List[int]], k: int, c: int) -> Tuple[List[List[int]], int]:
    """나무를 제거하고 제초제를 뿌립니다."""
    n = len(arr)
    max_remove_tree = 0
    max_remove_tree_y, max_remove_tree_x = -1, -1

    for i in range(n):
        for j in range(n):
            if arr[i][j] > 0:
                remove_tree_num = check_remove_tree(arr, i, j, k)
                if max_remove_tree < remove_tree_num:
                    max_remove_tree = remove_tree_num
                    max_remove_tree_y, max_remove_tree_x = i, j

    # 제초제 뿌리기
    arr[max_remove_tree_y][max_remove_tree_x] = -c
    for d in range(4):
        for mul in range(1, k+1):
            nx = max_remove_tree_x + DIAGONAL_DX[d] * mul
            ny = max_remove_tree_y + DIAGONAL_DY[d] * mul
            if not is_valid(nx, ny, n):
                break
            if arr[ny][nx] <= 0:
                if arr[ny][nx] > -c:
                    arr[ny][nx] = -c
                break
            arr[ny][nx] = -c

    return arr, max_remove_tree

def remove_poison(arr: List[List[int]]) -> List[List[int]]:
    """제초제 효과를 감소시킵니다."""
    n = len(arr)
    for i in range(n):
        for j in range(n):
            if arr[i][j] < 0:
                arr[i][j] += 1
    return arr

def main():
    """메인 함수: 시뮬레이션을 실행합니다."""
    n, m, k, c, arr = read_input()
    result = 0
    c += 1

    # 벽을 -99999로 설정
    for i in range(n):
        for j in range(n):
            if arr[i][j] == -1:
                arr[i][j] = -99999

    for _ in range(m):
        arr = grow_tree(arr)
        arr = reproduce_tree(arr)
        arr, remove_tree = exterminate_tree(arr, k, c)
        result += remove_tree
        arr = remove_poison(arr)

    print(result)

if __name__ == "__main__":
    main()