(Python) 백준 16235번: 나무 재테크

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

16235번: 나무 재테크

---

이 문제는 정말 여러번 풀어봤다.

문제 자체는 어렵지 않지만, 시간초과때문에 다소 빡셌다.

결국엔 시간때문에 다른 사람들의 코드를 많이 참고하게 되었던것같다.

 

 

# 1번 Code (python 시간초과, pypy3로 제출)

 

from collections import deque
 
n, m, k = map(int, input().split())
 
A = [list(map(int, input().split())) for _ in range(n)]
board = [[5] * n for _ in range(n)]
tree_dict = {(i, j): deque() for i in range(n) for j in range(n)}
 
dx = (-1, -1, 0, 1, 1, 1, 0, -1)
dy = (0, 1, 1, 1, 0, -1, -1, -1)
 
for _ in range(m):
    x, y, age = list(map(int, input().split()))
    tree_dict[(x - 1, y - 1)].append(age)
 
for year in range(k):
    # 봄
    for (x, y), trees in tree_dict.items():
        survive = deque()
        sumv = 0
        for tree in trees:
            if tree <= board[x][y]:
                board[x][y] -= tree
                survive.append(tree + 1)
            else:
                sumv += tree // 2
        tree_dict[(x, y)] = survive
        board[x][y] += sumv
    # 가을
    new_tree = []
    for (x, y), trees in tree_dict.items():
        for tree in trees:
            if tree % 5 == 0:
                for i in range(8):
                    nx, ny = x + dx[i], y + dy[i]
                    if not (0 <= nx < n and 0 <= ny < n): continue
                    new_tree.append((nx, ny))
    for x, y in new_tree:
        tree_dict[(x, y)].appendleft(1)
    # 겨울
    for i in range(n):
        for j in range(n):
            board[i][j] += A[i][j]
 
cnt = 0
for trees in tree_dict.values():
    cnt += len(trees)
print(cnt)

 

---

 

# 2번 Code (python 통과)

 

2번째 방법으로 풀면서 느낀건데, 파이썬은 딕셔너리를 잘 활용할수록 좋다!

n, m, k = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(n)]
board = [[5] * n for _ in range(n)]
tree = [[{} for _ in range(n)] for _ in range(n)]
 
dx = (-1, -1, -1, 0, 1, 1, 1, 0)
dy = (-1, 0, 1, 1, 1, 0, -1, -1)
 
for _ in range(m):
    x, y, age = map(int, input().split())
    tree[x - 1][y - 1][age] = 1
 
for year in range(k):
    # 봄, 여름
    for i in range(n):
        for j in range(n):
            if tree[i][j]:
                temp = {}
                die = 0
                for age in sorted(tree[i][j].keys()):
                    if age * tree[i][j][age] <= board[i][j]:
                        temp[age + 1] = tree[i][j][age]
                        board[i][j] -= age * tree[i][j][age]
                    else:
                        survive = board[i][j] // age
                        if not survive:  # 생존한 나무가 없음
                            die += (age // 2) * tree[i][j][age]
                            continue
                        board[i][j] -= age * survive
                        temp[age + 1] = survive
                        die += (age // 2) * (tree[i][j][age] - survive)
                tree[i][j] = temp
                board[i][j] += die
    # 가을
    for i in range(n):
        for j in range(n):
            if tree[i][j]:
                num = 0
                for age in tree[i][j].keys():
                    if age % 5 == 0:
                        num += tree[i][j][age]
                if num:
                    for d in range(8):
                        nx, ny = i + dx[d], j + dy[d]
                        if 0 <= nx < n and 0 <= ny < n:
                            if 1 not in tree[nx][ny].keys():
                                tree[nx][ny] = num
                            else:
                                tree[nx][ny] += num
    # 겨울
    for i in range(n):
        for j in range(n):
            board[i][j] += A[i][j]
cnt = 0
for i in range(n):
    for j in range(n):
        cnt += sum(tree[i][j].values())
print(cnt)