[코드트리/삼성SW역량테스트] 나무박멸 (Python)

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

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

 

Difficulty : Gold 4

 

Status : Solved

 

Time : 00:41:14

 

---

 

문제 설명 / 입력 및 출력

 

 

자세한 설명은 코드트리 사이트 링크를 참조해 주세요!

 

 

---

 

풀이

 

그나마 간단한(?) 구현 문제. 구현에 20여분, 디버깅에 20여분 소요되었다.

 

핵심은 "제초제가 어떻게 퍼지는지 정확하게 선지를 읽었는가?"가 되겠다. 제초제가 나무가 없는 곳에서는 퍼지지 않는다는 점, 나무가 있는 곳에서 퍼져도 나무가 없는 곳에 도달하면 그 지점까지만 탐색한다는 점을 고려해주어야 한다(나는 이 점에서 무작정 풀다가 추가시간이 소요되었다..)

 

 

풀이 코드

import sys
input = sys.stdin.readline

dr = [0, 0, -1, 1]
dc = [-1, 1, 0, 0]
ddr = [-1, -1, 1, 1]
ddc = [-1, 1, -1, 1]

n, m, K, C = map(int, input().split())

kills = [[0]*n for _ in range(n)]
maps = [list(map(int, input().split())) for _ in range(n)]
ans = 0

def growth() :
    for i in range(n) :
        for j in range(n) :
            if maps[i][j] > 0 :
                cnt = 0
                for k in range(4) :
                    ai, aj = i + dr[k], j + dc[k]
                    if -1 < ai < n and -1 < aj < n and maps[ai][aj] > 0 :
                        cnt += 1
                maps[i][j] += cnt

def breeding() :
    breeding_maps = [[0]*n for _ in range(n)]
    for i in range(n) :
        for j in range(n) :
            if maps[i][j] > 0 :
                breedings_list = list()
                for k in range(4) :
                    ai, aj = i + dr[k], j + dc[k]
                    if -1 < ai < n and -1 < aj < n and maps[ai][aj] == 0 and kills[ai][aj] == 0 :
                        breedings_list.append((ai, aj))
                if not breedings_list :
                    continue
                tree = maps[i][j] // len(breedings_list)
                for ai, aj in breedings_list :
                    breeding_maps[ai][aj] += tree
    
    for i in range(n) :
        for j in range(n) :
            maps[i][j] += breeding_maps[i][j]

def spread_test(i, j) :
    result = maps[i][j]
    if maps[i][j] == 0 :
        return 0
    for k in range(4) :
        for l in range(1, K+1) :
            ai, aj = i + ddr[k]*l, j + ddc[k]*l
            if not (-1 < ai < n and -1 < aj < n) or maps[ai][aj] in [0, -1]:
                break
            result += maps[ai][aj]
    return result

def spread(i, j) :
    kills[i][j] = C+1
    if maps[i][j] == 0 :
        return

    maps[i][j] = 0
    kills[i][j] = C+1
    for k in range(4) :
        for l in range(1, K+1) :
            ai, aj = i + ddr[k]*l, j + ddc[k]*l
            if not (-1 < ai < n and -1 < aj < n) :
                break
            if maps[ai][aj] in [0, -1] :
                kills[ai][aj] = C+1
                break
            maps[ai][aj] = 0
            kills[ai][aj] = C+1

def spread_spray() :
    global ans
    best_tree, bi, bj = 0, 0, 0
    for i in range(n) :
        for j in range(n) :
            if maps[i][j] == -1 :
                continue
            tmp = spread_test(i, j)
            if tmp > best_tree :
                best_tree, bi, bj = tmp, i, j
    if not best_tree :
        return False
    spread(bi, bj)
    ans += best_tree
    return True


def kills_update() :
    for i in range(n) :
        for j in range(n) :
            kills[i][j] = max(0, kills[i][j]-1)

for _ in range(m) :
    growth()
    breeding()
    is_tree = spread_spray()
    if not is_tree :
        break
    kills_update()

print(ans)

풀이 완료!