코드트리 나무박멸 (파이썬)

https://www.codetree.ai/training-field/frequent-problems/problems/tree-kill-all/

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

import sys

input = sys.stdin.readline

# 앞에 4개는 상하좌우방향, 뒤에 4개는 대각선방향
dx = [1, 0, -1, 0, 1, 1, -1, -1]
dy = [0, 1, 0, -1, 1, -1, 1, -1]

n, m, k, c = map(int, input().split())

a = [list(map(int, input().split())) for _ in range(n)]

ans = 0
kill_tree = [[0 for _ in range(n)] for _ in range(n)]  # 제초제 남은 기간 저장
for year in range(m):
    # 나무 성장
    for i in range(n):
        for j in range(n):
            if a[i][j] > 0:
                cnt = 0  # 인접한 나무 개수
                for d in range(4):
                    nx = i + dx[d]
                    ny = j + dy[d]
                    if 0 <= nx < n and 0 <= ny < n:
                        if a[nx][ny] > 0:
                            cnt += 1
                a[i][j] += cnt

    # 나무 번식
    # 중복 처리를 방지하기 위해서 add_tree에 번식된 나무만을 따로 저장한 후 나중에 갱신해줌
    add_tree = [[0 for _ in range(n)] for _ in range(n)]
    for i in range(n):
        for j in range(n):
            if a[i][j] > 0:
                cnt = 0  # 인접한 빈 칸 개수
                temp = []  # 빈 칸 좌표를 저장
                for d in range(4):
                    nx = i + dx[d]
                    ny = j + dy[d]
                    if 0 <= nx < n and 0 <= ny < n:
                        # 빈 칸 이고 제초제가 없어야함
                        if a[nx][ny] == 0 and kill_tree[nx][ny] == 0:
                            cnt += 1
                            temp.append([nx, ny])
                # 빈 칸이 있는 경우 add_tree에 번식된 나무만 저장
                if cnt > 0:
                    for x, y in temp:
                        add_tree[x][y] += a[i][j] // cnt
                        
    # 번식 진행
    for i in range(n):
        for j in range(n):
            if add_tree[i][j] > 0:
                a[i][j] = add_tree[i][j]

    # 나무 박멸
    # max_kill: 최대 박멸 그루 수
    max_kill = 0
    for i in range(n):
        for j in range(n):
            if a[i][j] > 0:
                cnt = a[i][j]  # 박멸 그루 수
                # stop_dir이 1이면 해당 방향은 더 진행안함
                stop_dir = [0 for _ in range(8)]
                for f in range(1, k+1):
                    for d in range(4, 8):
                        if stop_dir[d] == 0:
                            nx = i + f * dx[d]
                            ny = j + f * dy[d]
                            if 0 <= nx < n and 0 <= ny < n:
                                # 빈 칸이거나 벽이면 해당 방향은 이제 박멸하지 않음
                                if a[nx][ny] <= 0:
                                    stop_dir[d] = 1
                                else:
                                    cnt += a[nx][ny]
                            else:
                                stop_dir[d] = 1
                # 최대 박멸 그루 수를 갱신하고, 제초제 시작 좌표를 max_x, max_y에 저장
                if cnt > max_kill:
                    max_kill = cnt
                    max_x, max_y = i
                    
    # 박멸할 수 있는 나무가 있으면 박멸 진행
    # 마지막에 제초제가 있는 칸을 1 빼줄 것을 고려해서 kill_tree에 c + 1로 저장
    if max_kill > 0:
        x, y = max_x, max_y
        ans += max_kill
        a[x][y] = 0
        kill_tree[x][y] = c + 1
        stop_dir = [0 for _ in range(8)]
        for f in range(1, k+1):
            for d in range(4, 8):
                if stop_dir[d] == 0:
                    nx = x + f * dx[d]
                    ny = y + f * dy[d]
                    if 0 <= nx < n and 0 <= ny < n:
                        kill_tree[nx][ny] = c + 1
                        if a[nx][ny] <= 0:
                            stop_dir[d] = 1
                        else:
                            a[nx][ny] = 0
                    else:
                        stop_dir[d] = 1
            
        # 제초제 있는 칸 갱신
        for i in range(n):
            for j in range(n):
                if kill_tree[i][j] > 0:
                    kill_tree[i][j] -= 1

print(ans)