[알고리즘] 나무 재테크 - 구현

문제

https://www.acmicpc.net/problem/16235

16235번: 나무 재테크

 

풀이 팁

1. 이 문제는 단순히 조건을 따라가면서 풀면 되는 문제인데, 가장 중요한 점은 시간복잡도 이다.

-> 나는 처음에 단순히 3차원 배열 만들면서 배열을 초기화 시키면서 했다 ( 가장 아래에 처음 코드를 넣어 놓겠다 )

-> 읽어보면 초기화를 엄청 시키면서 했는데, 나는 이 부분이 시간초과의 원인이라고 생각한다 ( 테코는 다 맞음 )

 

2. 시간 초과를 줄이기 위해서 한일 ( 여기 팁은 질문 게시판에서 얻었다 )

- 3차원 배열의 가장 안쪽은 list가 아닌 deque로 접근

- sys.stdin.readline 사용 ( 근데 이건 크게 도움이 된지 모르겠음 )

- 함수 줄이기 ( spring, summer, fall, winter 4개를 중복되는 코드는 묶어서 2개로 줄임 )

 

코드

import sys
from collections import deque

input = sys.stdin.readline

n,m,k = map(int,input().split())

dxs, dys = [-1,-1,-1,0,1,1,1,0],[-1,0,1,1,1,0,-1,-1]

# 겨울마다 주는 양분
extra_feed = [
    list(map(int,input().split()))
    for _ in range(n)
]

# 트리
tree = [
    [deque() for _ in range(n)]
    for _ in range(n)
]

# 기본 양분
feed = [
    [5 for _ in range(n)]
    for _ in range(n)
]

for _ in range(m):
    x,y,age = map(int,input().split())
    tree[x-1][y-1].append(age)

def in_range(x,y):
    return 0<=x<n and 0<=y<n

def spring_summer():
    for i in range(n):
        for j in range(n):
            dead = 0
            new_tree = deque()
            for age in tree[i][j]:

                if feed[i][j] >= age:
                    feed[i][j] -= age
                    new_tree.append(age+1)
                else:
                    dead += age // 2
            
            feed[i][j] += dead
            tree[i][j] = new_tree

def fall_winter():
    tmp_tree = []
    for x in range(n):
        for y in range(n):
            for age in tree[x][y]:
                if age % 5 == 0:
                    for dx, dy in zip(dxs, dys):
                        nx,ny = x + dx, y + dy

                        if not in_range(nx,ny):
                            continue

                        tmp_tree.append((nx,ny))
                else:
                    continue
            feed[x][y]  += extra_feed[x][y]
            
    for pos in tmp_tree:
        x,y = pos
        tree[x][y].appendleft(1)

ans = 0
def count_tree():
    global ans

    for i in range(n):
        for j in range(n):
            ans += len(tree[i][j])

for _ in range(k):
    spring_summer()
    fall_winter()

count_tree()
print(ans)

 

 

 

[ 처음 코드 ( 참조용 ) ]

import sys
from collections import deque
import heapq

input = sys.stdin.readline

n,m,k = map(int,input().split())

dxs, dys = [-1,-1,-1,0,1,1,1,0],[-1,0,1,1,1,0,-1,-1]

# 겨울마다 주는 양분
extra_feed = [
    list(map(int,input().split()))
    for _ in range(n)
]

# 트리
tree = [
    [[] for _ in range(n)]
    for _ in range(n)
]

temp_tree = [
    [[] for _ in range(n)]
    for _ in range(n)
]

dead_tree = []

# 기본 양분
feed = [
    [5 for _ in range(n)]
    for _ in range(n)
]

for _ in range(m):
    x,y,age = map(int,input().split())
    tree[x-1][y-1].append(age)

def in_range(x,y):
    return 0<=x<n and 0<=y<n

def init_tree():
    for i in range(n):
        for j in range(n):
            temp_tree[i][j] = []

def clear_dead_tree():
    global dead_tree 

    dead_tree = []

def spring():
    global dead_tree
    init_tree()

    for i in range(n):
        for j in range(n):
            tree[i][j].sort()
            
            idx = 0

            while feed[i][j] > 0:
                if idx >= len(tree[i][j]):
                    break

                if feed[i][j] >= tree[i][j][idx]:
                    feed[i][j] -= tree[i][j][idx]
                    idx += 1
                else:
                    feed[i][j] = 0
                    break
            
            for age in tree[i][j][idx:]:
                dead_tree.append((i,j,age))

            tree[i][j] = tree[i][j][:idx] 
    
    
    for i in range(n):
        for j in range(n):
            for age in tree[i][j]:
                temp_tree[i][j].append(age+1)
    
    for i in range(n):
        for j in range(n):
            tree[i][j]  = temp_tree[i][j]

def summer():
    global dead_tree
    
    for i,j,age in dead_tree:
        feed[i][j] += age // 2

def fall_winter():
    for x in range(n):
        for y in range(n):
            for age in tree[x][y]:
                if age % 5 == 0:
                    for dx,dy in zip(dxs, dys):
                        nx,ny = x + dx, y + dy

                        if not in_range(nx,ny):
                            continue

                        tree[nx][ny].append(1)
                else:
                    continue
            
            feed[x][y] += extra_feed[x][y]

ans = 0
def count_tree():
    global ans

    for i in range(n):
        for j in range(n):
            ans += len(tree[i][j])

for _ in range(k):
    spring()
    summer()
    fall_winter()

count_tree()
print(ans)