[백준] 16235. 나무 재테크(python, 파이썬)

문제

16235. 나무 재테크

📌 아이디어

봄, 여름, 가을, 겨울에 맞는 각각의 조건에 맞춰서 구현하면 되는 문제였습니다.

• 봄 : 나무가 1개 이상 존재하는 경우는 나무 나이별로 오름차순으로 정렬하여 양분을 부여합니다. 하지만 나이보다 양분이 작다면 그 이후로는 양분을 배분할 수 없기에 종료시점의 인덱스를 저장하고 반복을 종료합니다.
• 여름 : 죽은 나무가 있다면 종료 시점까지의 길이 만큼 나무들을 슬라이싱합니다. 그리고 죽은나무들은 양분으로 배분됩니다.
• 가을 : 나무가 1개 이상 존재하는 경우에 5의 배수인지를 확인하여 8방향으로 나무를 번식합니다.
• 겨울 : input 받은 양분을 배분합니다.

📌 코드

import sys
input = sys.stdin.readline
N, M, K = map(int, input().split())
add = [list(map(int, input().split())) for _ in range(N)]
info = [list(map(int, input().split())) for _ in range(M)]

age = [[[] for _ in range(N)] for _ in range(N)]
pos = [[5] * N for _ in range(N)] # 양분 저장

for x, y, z in info:
    age[x-1][y-1].append(z) # 나이

for _ in range(K):
    # spring
    for i in range(N):
        for j in range(N):
            death_note = -1
            if len(age[i][j]) >= 1:
                age[i][j].sort()
                for k in range(len(age[i][j])):
                    if pos[i][j] < age[i][j][k]:
                        # 종료 시점 저장
                        death_note = k
                        break
                    else:
                        # 양분 배분 및 나이 증가
                        pos[i][j] -= age[i][j][k]
                        age[i][j][k] += 1

            # summer
            if death_note != -1:
                length = len(age[i][j])
                tmp = age[i][j]
                # 종료시점까지 나무를 슬라이싱
                age[i][j] = age[i][j][:death_note]
                for k in range(death_note, length):
                    # 양분 추가
                    pos[i][j] += (tmp[k] // 2)

    # fall
    for i in range(N):
        for j in range(N):
            if len(age[i][j]) >= 1:
                for k in range(len(age[i][j])):
                    if age[i][j][k] >= 5 and age[i][j][k] % 5 == 0:
                        for di, dj in ((1, 0), (0, 1), (-1, 0), (0, -1), (1, 1), (-1, -1), (1, -1), (-1, 1)):
                            ni = i + di
                            nj = j + dj
                            # 나무 번식
                            if 0 <= ni < N and 0 <= nj < N:
                                age[ni][nj].append(1)
    # winter
    for i in range(N):
        for j in range(N):
            # 양분 추가
            pos[i][j] += add[i][j]

ans = 0
for i in range(N):
    for j in range(N):
        if len(age[i][j]) >= 1:
            ans += len(age[i][j])

print(ans)

피드백 및 개선점은 댓글을 통해 알려주세요😊