[삼성코테/파이썬] 나무박멸

https://www.codetree.ai/problems/tree-kill-all/submissions

코드 길이가 100줄이 넘기 때문에 자잘한 실수라도 있으면 틀릴 수 밖에 없는 문제다. 당연히 한 번에 푼 것은 아니고 각 함수들의 뼈대를 먼저 구성한 다음에 구체적인 요건들을 조금씩 더해가면서 풀었다. 그럼에도 불구하고 오답이었는데, 그 이유는 최대 박멸 장소를 찾는 def check 함수 구현에 return 값을 list로 만들었기 때문이다. return 값이 존재하지 않을 수 있다는 사실을 추가적으로 리서치해보고 return 값을 list가 아닌 default int 값으로 바꾸어주었다. 그렇게 되면 가장 앞에 오는 최대 박멸 장소가 정해지면 그 이후로는 update 되지 않는다. (우선순위는 따라서 열-행 순서로 for 순회를 만들었다)

그 밖에도 사소한 실수들이 있었는데, 제초제를 뿌릴 때 장애물이나 빈 땅이 있으면 그 이후로 진행되지 말아야 하는데, for 순회의 순서(방향-깊이)를 처음에 잘 못 짜기도 했고 break 조건이 정확하지 않아서 오답이 나왔다.

무엇보다도 구현이므로 시간이 엄청 오래 걸렸다.

N, M, K, C = map(int, input().split())
from collections import defaultdict
#2차원의 상태 [나무, 제초제 만기 시간] 혹은 각각의 1차원 array를 총 2개 가지고 한다 (후자 선택)
arr = [[*map(int,input().split())] for _ in range(N)]
ptime = [[0]*N for _ in range(N)]
dy = [0,1,0,-1]
dx = [1,0,-1,0]
answer = 0
# 성장

def count_something(x, y, some):
    cnt = 0
    for i in range(4):
        ny = y + dy[i]
        nx = x + dx[i]
        if ny < 0 or nx < 0 or ny > N - 1 or nx > N - 1 or ptime[ny][nx] or arr[ny][nx] == -1:
            continue
        if some == 'free':
            if arr[ny][nx] == 0:
                cnt += 1
        elif some == 'tree':
            if arr[ny][nx] > 0:
                cnt +=1
    return cnt


def grow(arr):

    for y in range(N):
        for x in range(N):
            if ptime[y][x] or arr[y][x] <= 0:
                continue
            arr[y][x] += count_something(x,y,'tree') # todo: 주변의 나무 개수를 count 하여 더해준다.
    return arr

def expansion(arr):
    narr = [[0] * N for _ in range(N)]
    for y in range(N):
        for x in range(N):
            if ptime[y][x]:
                continue
            free = count_something(x,y,'free')
            if not free:
                continue
            rep = arr[y][x] // free #todo: 실제 주변에 번식 가능여부를 cnt 한다.
            if rep > 0:
                for i in range(4):
                    ny = y + dy[i]
                    nx = x + dx[i]
                    if ny < 0 or nx < 0 or ny > N-1 or nx > N-1  or ptime[ny][nx] or arr[ny][nx] > 0 or arr[ny][nx] == -1:
                        continue
                    else:
                        narr[ny][nx] += rep

    for y in range(N):
        for x in range(N):
            if ptime[y][x]:
                continue
            narr[y][x] += arr[y][x]

    return narr

dy2 = [1,-1,-1,1] # 대각선 움직임
dx2 = [1,-1,1,-1]


def reset_pest():
    for y in range(N):
        for x in range(N):
            if ptime[y][x]:
                ptime[y][x] -= 1

def check(arr,k)->list:
    global answer
    mx = 0
    mx_x = 1
    mx_y = 1
    d = defaultdict(list)
    for y in range(N):
        for x in range(N):
            if ptime[y][x] or arr[y][x] <= 0:
                continue
            cut = 0
            cut += arr[y][x]
            for i in range(4):
                for kk in range(1, k + 1):
                    ny = y + dy2[i] * kk
                    nx = x + dx2[i] * kk
                    if ny < 0 or nx < 0 or ny > N - 1 or nx > N - 1 or ptime[ny][nx] or arr[ny][nx] <= 0:
                        break
                    cut+=arr[ny][nx]
            if mx < cut:
                mx = cut
                mx_x = x
                mx_y = y

    return mx, mx_x, mx_y

def eliminate(arr,k):
    global answer

    count, x,y = check(arr,k) # 제초제를 뿌릴 위치를 찾는다.
    if not count:
        return arr # 종료
    answer += count

    ptime[y][x] = C
    for i in range(4):
        for kk in range(1, k + 1):
            ny = y + dy2[i]*kk
            nx = x + dx2[i]*kk
            if ny < 0 or nx < 0 or ny > N - 1 or nx > N - 1 or arr[ny][nx] == -1:
                break
            if arr[ny][nx] == 0:
                ptime[ny][nx] = C
                break
            ptime[ny][nx] = C
    return arr


for m in range(M):

    arr = grow(arr) # 성장한다.
    arr = expansion(arr) # 번식한다.
    reset_pest() # pest 의 수명을 -1 감소시킨다.
    arr = eliminate(arr,K) # 제초제를 뿌린다.

print(answer)