코드트리 포탑 부수기 (파이썬)

https://www.codetree.ai/training-field/frequent-problems/problems/destroy-the-turret

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

import sys
from collections import deque

input = sys.stdin.readline

# 레이저 공격 우하좌상 순 + 포탄공격 대각선
dx = [0, 1, 0, -1, 1, 1, -1, -1]
dy = [1, 0, -1, 0, 1, -1, 1, -1]

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

# attack에 공격한 시점 저장
attack = [[0 for _ in range(m)] for _ in range(n)]
for kk in range(k):
    # cand에 부서지지 않은 모든 포탑을 저장
    cand = []
    for i in range(n):
        for j in range(m):
            if a[i][j] > 0:
                cand.append([a[i][j], attack[i][j], i, j])

    # 부서지지 않은 포탑이 한개면 종료
    if len(cand) <= 1:
        break

    # 문제의 조건대로 정렬 후 공격자와 공격을 받는 포탑 선정
    cand.sort(key=lambda x: (x[0], -x[1], -(x[2] + x[3]), -x[3]))
    sx, sy = cand[0][2], cand[0][3]
    attack[sx][sy] = kk + 1
    tx, ty = cand[-1][2], cand[-1][3]
    a[sx][sy] += n + m

    # bfs로 레이저 공격이 가능한 최단경로탐색
    q = deque()
    q.append([sx, sy])
    check = [[0 for _ in range(m)] for _ in range(n)]
    dmg = [[0 for _ in range(m)] for _ in range(n)]
    # dmg는 공격에 참여한 포탑의 좌표를 1로 해서 포탑 정비때 사용
    dmg[sx][sy], dmg[tx][ty] = 1, 1
    check[sx][sy] = 1
    flag = 0
    while q:
        x, y = q.popleft()
        # 공격 받는 포탑의 좌표가 나오면 레이저 공격 실행
        if x == tx and y == ty:
            flag = 1
            a[x][y] -= a[sx][sy]
            if a[x][y] < 0:
                a[x][y] = 0
            while True:
                x, y = check[x][y]
                if x == sx and y == sy:
                    break
                a[x][y] -= a[sx][sy] // 2
                dmg[x][y] = 1
                if a[x][y] < 0:
                    a[x][y] = 0

        if flag == 1:
            break

        for i in range(4):
            nx = (x + dx[i]) % n
            ny = (y + dy[i]) % m
            if check[nx][ny] == 0 and a[nx][ny] > 0:
                # check에는 이전 좌표를 저장해서 이동 경로를 계속 기록
                check[nx][ny] = [x, y]
                q.append([nx, ny])

    # 레이저 공격이 불가능하면 포탄 공격
    if flag == 0:
        a[tx][ty] -= a[sx][sy]
        if a[tx][ty] < 0:
            a[tx][ty] = 0
        for i in range(8):
            nx = (tx + dx[i]) % n
            ny = (ty + dy[i]) % m
            if [nx, ny] != [sx, sy]:
                if a[nx][ny] > 0:
                    a[nx][ny] -= a[sx][sy] // 2
                    dmg[nx][ny] = 1
                    if a[nx][ny] < 0:
                        a[nx][ny] = 0

    # 포탑 정비
    for i in range(n):
        for j in range(m):
            if a[i][j] > 0 and dmg[i][j] == 0:
                a[i][j] += 1

print(max(map(max, a)))