[Python/Codetree] 생명과학부 랩 인턴

https://www.codetree.ai/training-field/frequent-problems/problems/biology-lab-intern/description?page=2&pageSize=20

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

문제

생명과학부 랩 인턴을 시작하게 된 승용이는 n * m 격자판에서 움직이는 곰팡이를 채취하는 일을 맡았습니다. 격자판에서 곰팡이를 채취할 때는 다음과 같은 규칙을 이용합니다. 아래 그림에서 빨간색으로 표시된 숫자는 곰팡이의 크기를 의미하고, 파란색으로 표시된 숫자는 속력을 의미합니다.

1. 승용이는 첫번째 열부터 탐색을 시작합니다.
2. 해당 열의 위에서 아래로 내려가며 탐색할 때 제일 빨리 발견한 곰팡이를 채취합니다. 곰팡이를 채취하고 나면 해당 칸은 빈칸이 되며, 해당 열에서 채취할 수 있는 곰팡이가 없는 경우도 있을 수 있음에 유의합니다.
3. 해당 열에서 채취 시도가 끝나고 나면 곰팡이는 이동을 시작합니다.
4. 입력으로 주어진 방향과 속력으로 이동하며 격자판의 벽에 도달하면 반대로 방향을 바꾸고 속력을 유지한 채로 이동합니다. 방향을 바꿀 때는 시간이 걸리지 않습니다.
5. 모든 곰팡이가 이동을 끝낸 후에 한 칸에 곰팡이가 두마리 이상일 때는 크기가 큰 곰팡이가 다른 곰팡이를 모두 잡아먹습니다.
6. 이 모든 과정은 1초가 걸리며 이후 승용이는 오른쪽 열로 이동해서 위의 과정을 반복합니다.

예를 들어 위의 상태에서 채취를 하고 난 뒤의 상태는 아래와 같이 변하게 됩니다.

이후에 곰팡이가 본인의 속도대로 이동하면 아래와 같은 상태로 변하게 됩니다.

끝나고 난 뒤 승용이는 열을 오른쪽으로 이동합니다.

이동 후에 5번째 행, 3번째 열에서 두 곰팡이가 부딪히게 됩니다. 이때 크기가 5인 곰팡이가 크기가 더 크기 때문에 나머지 곰팡이를 모두 잡아먹습니다. 그리고 나서 아래와 같이 바뀌게 됩니다.

승용이가 해당 격자판에 있는 모든 열을 검사했을 때, 채취한 곰팡이 크기의 총합을 구하는 프로그램을 구해보세요.

풀이

Code

from collections import defaultdict

# 1. 곰팡이 제거 함수 정의
def remove(col) :
    global ans

    # 1-1. 해당 열에 포함되는 곰팡이 위치 정렬
    index = sorted([(x, y) for (x, y) in information.keys() if y == col])
    # 1-2. 곰팡이가 있다면 첫 곰팡이 크기 카운팅 후 제거
    if index :
        ans += information[index[0]][-1]
        information.pop(index[0])

# 2. 곰팡이 이동 함수 정의
def move() :
    # 2-1. 곰팡이 정보 저장 함수 정의
    def update(x, y, s, d, b) :
        # 2-1-1. 해당 위치가 비었다면 정보 저장
        if not updated_information[(x, y)] : updated_information[(x, y)] = [s, d, b]
        # 2-1-2. 해당 위치가 비어있지 않을 경우
        else :
            # 현재 입력받은 곰팡이 크기가 더 클 경우 업데이트
            if updated_information[(x, y)][-1] < b :
                updated_information[(x, y)] = [s, d, b]
    # 2-2. 반환용 딕셔너리 생성
    updated_information = defaultdict(list)
    # 2-3.
    for key in information.keys() :
        # 2-3-1. 정보 반환
        x, y = key
        s, d, b = information[key]
        # 2-3-2. 속력이 0인 경우
        if s == 0 : update(x, y, s, d, b)
        nx, ny = s * dirs[d][0] + x, s * dirs[d][1] + y
        # 2-3-3. 격자 칸에 부딪히지 않는 경우
        if nx > 0 and nx <= n and ny > 0 and ny <= m :
            # 끝에 도달했다면 방향만 변화
            if (d == 1 and nx == 1) or (d == 2 and nx == n) or (d == 3 and ny == m) or (d == 4 and ny == 1) : d = now_dir[d]
            update(nx, ny, s, d, b)
        # 2-3-4. 격자칸을 넘어가는 경우
        else :
            # 남은 거리 계산
            if d == 1 : mod = abs(nx) + 1
            elif d == 2 : mod = nx - n
            elif d == 3 : mod = ny - m
            else : mod = abs(ny) + 1
            # 갈 수 있는 칸
            if d in [1, 2] : ability = n - 1
            else : ability = m - 1
            # (남은 거리 // 갈 수 있는 칸) 값이 짝수일 경우
            if (mod // ability) % 2 == 0 :
                # 현재와 반대 방향으로 전환
                d = now_dir[d]
            # 다음 위치 설정
            mod %= ability
            if d == 1 : nx, ny = n + dirs[d][0] * mod, y
            elif d == 2 : nx, ny = 1 + dirs[d][0] * mod, y
            elif d == 3 : nx, ny = x, 1 + dirs[d][1] * mod
            else : nx, ny = x, m + dirs[d][1] * mod
            update(nx, ny, s, d, b)
    # 2-4. 딕셔너리 반환
    return updated_information

ans = 0
n, m, k = map(int, input().split())
dirs = [[], (-1, 0), (1, 0), (0, 1), (0, -1)]
now_dir = {1 : 2, 2 : 1, 3 : 4, 4 : 3}
# 3. 곰팡이 정도 딕셔너리로 입력받기
information = defaultdict(list)
for _ in range(k) :
    x, y, s, d, b = map(int, input().split())
    information[(x, y)] = [s, d, b]
# 4.
for i in range(1, m+1) :
    # 4-1. 곰팡이 채취
    remove(i)
    # 4-2. 곰팡이 이동
    information = move()
# 5. 결과 출력
print(ans)