[백준 17135] 캐슬 디펜스(Python)

4개월 동안 방치한 문제를 이제야 풀었다.

1. 요약

N*M 격자판 어딘가에 적이 있고, N+1행, 모든 칸에 성이 있다. 캐슬 디펜스는 궁수 3명을 성에 배치 할 때, 주어진 공격 범위안의 가장 가까운 적 1명을 퇴치하는 게임이다. 적이 성이 있는 칸으로 이동한 경우 게임에서 제외되며, 모든 적이 격자판에서 제외되어야 게임이 끝난다. 이 때, 3명의 궁수를 성에 배치해서 잡을 수 있는 최대 적의 수를 구해야한다.(단 사거리는 |y1-y2|+|x1-x2|로 정의되며, 궁수와 가장 가까운 적이 여러개이면, 가장 왼쪽을 먼저 공격한다.)

2. 아이디어

3명의 궁수를 [0,1,2,3..m-1]에 배치 할 때, 나올 수 있는 경우의 수는 5C3 = 10가지이다. 다시 말해, 조합을 이용해서, 궁수 3명 배치하여 나올 수 있는 모든 경우를 구한다. 그리고 각 경우에 대해, 잡을 수 있는 적들의 수를 구한다. 여기서 가장 큰 값이 정답이다. 가장 가까운 적이 다수면 가장 왼쪽에 있는 적을 공격한다고 문제에 언급됐는데, 이는 우선순위 큐로 구현한다.

3. 코드

import sys
from itertools import combinations
from heapq import heappop, heappush

n, m, d = map(int, sys.stdin.readline().rstrip().split())
MAP = [list(map(int, sys.stdin.readline().rstrip().split())) for _ in range(n)]
arr = []

castle_pos = [i for i in range(m)]
archer_cases = tuple(combinations(castle_pos, 3))
answer = 0

def get_enemy_count():
    global arr
    cnt = 0
    
    for i in range(n):
        for j in range(m):
            if arr[i][j] == 1: cnt += 1
    
    return cnt

def attack_enemy(archer_case_index):
    global arr
    remove_list = []
    attacked = [[False for _ in range(m)] for _ in range(n)]
    cnt = 0

    for archer_pos in archer_cases[archer_case_index]:
        pq = [] # [거리, x, y]를 우선순위 큐에 삽입

        for i in range(n-1, -1, -1):
            for j in range(m):
                if arr[i][j] == 1:
                    diff = abs(n-i) + abs(archer_pos-j)
                    if diff <= d:
                        heappush(pq, [diff, j, i])

        if pq:  
            _, x, y = heappop(pq)
            remove_list.append([y, x])
    
    for y, x in remove_list:
        if not attacked[y][x]:
            attacked[y][x] = True
            cnt += 1
            arr[y][x] = 0
    
    return cnt

def move_enemy():
    global arr
    arr[-1] = [0 for _ in range(m)]

    for i in range(-1, -n, -1):
        arr[i] = arr[i-1]
    
    arr[0] = [0 for _ in range(m)]

def simulation(archer_case_index):
    cnt = 0

    while get_enemy_count() != 0:
        cnt += attack_enemy(archer_case_index)
        move_enemy()
        
    return cnt

for i in range(len(archer_cases)):
    arr = [[MAP[i][j] for j in range(m)] for i in range(n)]
    cnt = simulation(i)
    if answer < cnt: answer = cnt

print(answer)