Jae-Kyung Cho LLM Developers who was a Robotics engineer

Baekjoon-16235-나무 재테크

from collections import deque
import sys
input = sys.stdin.readline

N,M,K = map(int,input().split())
A = [list(map(int,input().split())) for _ in range(N)]
trees = [list(map(int,input().split())) for _ in range(M)]

class Farm:
    def __init__(self, trees, A):
        self.trees = self.initialize(trees)
        self.A = A
        self.feed = [[5]*N for _ in range(N)]
    
    def initialize(self, trees):
        tree_dict = [[deque() for _ in range(N)] for _ in range(N)]
        for x,y,age in trees:
            tree_dict[x-1][y-1].append(age)
        return tree_dict

    def season(self):
        # print('---1-----')
        # print(*self.trees,sep='\n')
        self.ss()
        self.autumn()
        self.winter()
        # print('----2----')
        # print(*self.trees,sep='\n')

    def count_trees(self):
        result = 0
        for x in range(N):
            for y in range(N):
                result += len(self.trees[x][y])
        return result

    def ss(self):
        for x in range(N):
            for y in range(N):
                tree_list = self.trees[x][y]
                new_list = deque()
                while len(tree_list):
                    age = tree_list.popleft()
                    if self.feed[x][y] >= age:
                        self.feed[x][y] -= age
                        new_list.append(age+1)
                    else:
                        tree_list.append(age)
                        break
                self.feed[x][y] += sum([a//2 for a in tree_list])
                self.trees[x][y] = new_list
    
    def autumn(self):
        for x in range(N):
            for y in range(N):
                for age in self.trees[x][y]:
                    if age%5 == 0:
                        for nx,ny in [[x-1,y-1],[x-1,y],[x-1,y+1],[x,y-1],[x,y+1],[x+1,y-1],[x+1,y],[x+1,y+1]]:
                            if 0<=nx<N and 0<=ny<N:
                                self.trees[nx][ny].appendleft(1)
    
    def winter(self):
        for i in range(N):
            for j in range(N):
                self.feed[i][j] += self.A[i][j]

farm = Farm(trees, A)
for _ in range(K):
    farm.season()
print(farm.count_trees())
Comments