[Python/Codetree] 나무 타이쿤

https://www.codetree.ai/training-field/frequent-problems/problems/tree-tycoon/submissions?page=1&pageSize=20&order=tier 

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

문제

신비한 나무 리브로수를 n x n 격자 칸에다가 키우고자 합니다.

n x n 격자에 서로 다른 높이를 가진 리브로수가 주어집니다.

리브로수를 키우기 위해서는 특수 영양제를 필요로 합니다. 특수 영양제는 1 x 1 땅에 있는 리브로수의 높이를 1 증가시키며, 만약 해당 땅에 씨앗만 있는 경우에는 높이 1의 리브로수를 만들어냅니다.

초기 특수 영양제는 n x n 격자의 좌하단의 4개의 칸에 주어집니다.

각각의 특수 영양제는 이동 규칙에 따라 움직이는데, 이동 규칙은 이동 방향과 이동 칸 수로 주어집니다. 이동 방향의 경우 1번부터 8번까지 → ↗ ↑ ↖ ← ↙ ↓ ↘으로 주어지며 이동 칸 수만큼 특수 영양제가 이동하게 됩니다. 격자의 모든 행,열은 각각 끝과 끝이 연결되어 있습니다. 즉 격자 바깥으로 나가면 마치 지구가 둥근것처럼 반대편으로 돌아옵니다. 만약 n번 열에서 오른쪽에서 이동하는 경우에는 1번 열으로 이동하게 됩니다.

1년동안 다음 단계를 거쳐 리브로수는 성장하게 됩니다.

1. 특수 영양제를 이동 규칙에 따라 이동시킵니다.
2. 특수 영양제를 이동 시킨 후 해당 땅에 특수 영양제를 투입합니다.
3. 특수 영양제를 투입한 리브로수의 대각선으로 인접한 방향에 높이가 1 이상인 리브로수가 있는 만큼 높이가 더 성장합니다. 대각선으로 인접한 방향이 격자를 벗어나는 경우에는 세지 않습니다.
4. 특수 영양제를 투입한 리브로수를 제외하고 높이가 2 이상인 리브로수는 높이 2를 베어서 잘라낸 리브로수로 특수 영양제를 사고, 해당 위치에 특수 영양제를 올려둡니다.

아래 예시를 통해서 1년 동안 리브로수의 성장이 어떻게 이뤄지는지 살펴보겠습니다.

다음과 같이 각기 높이가 다른 리브로수가 주어질 때(높이가 0인 리브로수는 새싹만 존재하는 땅을 의미합니다.), 초기의 특수 영양제의 위치는 아래 음영이 있는 칸이 됩니다. 이때 주어진 이동 규칙에 따라 이동하면 아래의 그림처럼 변합니다.

이때 특수 영양제가 있는 땅의 리브로수는 높이가 1만큼 증가하고 대각선으로 인접한 높이 1 이상의 리브로수의 개수 만큼 높이가 증가합니다.

이후 해당 년도에 특수 영양제를 맞은 땅을 제외하고 높이가 2 이상인 리브로수를 높이 2만큼 잘라내고 해당 땅 위에 특수 영양제를 올려줍니다.

다음 년도로 넘어가면 위의 4단계를 다시 반복합니다.
n x n 격자의 칸에 리브로수의 각자 다른 높이와 각 년도의 이동 규칙이 주어질 때, 해당 년수가 모두 지나고 난 뒤 남아있는 리브로수 높이들의 총 합을 구하는 프로그램을 작성하세요.

입력 형식

첫번째 줄에는 격자의 크기 n, 리브로수를 키우는 총 년 수 m이 주어집니다.

이후 두번째 줄부터 n+1번째 줄까지 서로 다른 리브로수의 높이가 주어집니다.

이후 m개의 줄에는 각 년도의 이동 규칙이 주어집니다.

이동 규칙은 이동 방향 d, 이동 칸 수 p로 주어지고, d는 1번부터 8번까지 각각 → ↗ ↑ ↖ ← ↙ ↓ ↘으로 주어집니다.

• 3 ≤ n ≤ 15
• 1 ≤ m ≤ 100
• 0 ≤ 초기에 주어지는 리브로수의 높이 ≤ 100
• 1 ≤ d ≤ 8
• 1 ≤ p ≤ min(n, 10)

출력 형식

m년 이후 남아있는 리브로수의 총 높이의 합을 출력하세요.

풀이

Code

import sys
input = sys.stdin.readline

def solution(n, m, graph, years) :
    # 1. 특수 영양제 이동 함수 정의
    def move(dir, how) :
        # 1-1.
        for idx in range(len(nutrients)) :
            # 위치 이동
            x, y = nutrients[idx][0], nutrients[idx][1]
            nx, ny = (x + dirs[dir][0] * how) % n, (y + dirs[dir][1] * how) % n
            nutrients[idx] = [nx, ny]

    # 2. 리브로수 성장 함수 정의
    def grow_up() :
        # 2-1. 반환용 graph 생성
        return_graph = graph.copy()
        # 2-2. 현재 영양제 위치의 리브로수 높이 1 성장
        for x, y in nutrients :
            graph[x][y] += 1
        # 2-3.
        for x, y in nutrients :
            count = 0
            for dir_x, dir_y in [(-1, -1), (-1, 1), (1, -1), (1, 1)] :
                nx, ny = x + dir_x, y + dir_y
                # 예외 처리
                if nx < 0 or nx >= n or ny < 0 or ny >= n or not graph[nx][ny] : continue
                # 카운트
                count += 1
            return_graph[x][y] += count
        # 2-4. graph 반환
        return return_graph

    # 3. 영양제 재배치 함수 정의
    def relocation() :
        # 3-1. 반환용 영양제 리스트 생성
        new_nutrients = []
        # 3-2.
        for i in range(n) :
            for j in range(n) :
                # 해당 위치가 이전 특수 영양제 위치에 속하지 않은 경우
                if [i, j] not in nutrients :
                    # 해당 위치의 리브로수 높이가 2 이상인 경우
                    if graph[i][j] >= 2 :
                        # 반환용 영양제 리스트에 인덱스 삽입
                        new_nutrients.append([i, j])
                        # 리브로수 높이 감소
                        graph[i][j] -= 2
        # 3-3. 새로운 영양제 리스트 반환
        return new_nutrients

    # 4. 방향 정보 리스트 생성
    dirs = [[], (0, 1), (-1, 1), (-1, 0), (-1, -1), (0, -1), (1, -1), (1, 0), (1, 1)]
    # 5. 초기 영양제 리스트 생성
    nutrients = [[n-2, 0], [n-2, 1], [n-1, 0], [n-1, 1]]
    # 6. 
    for dir, how in years :
        # 특수 영양제 이동
        move(dir, how)
        # 리브로수 성장
        graph = grow_up()
        # 특수 영양제 재배치
        nutrients = relocation()
    # 7. 리브로수의 총 높이 출력
    print(sum([sum(line) for line in graph]))

if __name__ == "__main__" :
    n, m = map(int, input().split())
    graph = [list(map(int, input().split())) for _ in range(n)]
    years = [list(map(int, input().split())) for _ in range(m)]
    solution(n, m, graph, years)