[BOJ] 20056. 마법사 상어와 파이어볼 (🥇 , 구현/시뮬레이션)

문제

풀이

• 너무 뻔한 구현문제라 어렵지 않았다. 문제가 제시하는 순서대로만 풀면 마법처럼 풀림~~~
• 다만 파이어볼이 2개 이상인 경우 라는 조건은 쉽게 낚일 수 있는 듯

수정 전 코드

• 2개 이상의 파이어볼이 들어 있는 경우 방향을 변경하게 되는 코드를 따로 함수를 빼지 않은 상태
• 파이어볼이 1개만 있는 경우를 if~else 문의 else 부분에 넣어놨는데, 조금 가독성이 떨어지는 느낌? 어디서 이 처리를 하고 있는지 눈에 잘 안 띄는 것 같다
• global 선언해서 사용하지 않으려고 계속 return 값으로 넘겨주는 방식으로 작성했는데 오히려 코드가 더 복잡해진 것 같다. 애초에 하나의 데이터를 가지고 계속 변경해가며 사용하는 거니까 전역 변수를 가져와서 사용하는 게 맞는 것 같당ㅇ..

import sys
input = sys.stdin.readline

def sol():

    N, M, K = map(int, input().split())
    DIR = {
        0: (-1, 0),
        1: (-1, 1),
        2: (0, 1),
        3: (1, 1),
        4: (1, 0),
        5: (1, -1),
        6: (0, -1),
        7: (-1, -1)
    }

    # 파이어볼 딕셔너리로 입력받기
    def INPUT(M):
        fireball = {}
        for _ in range(M):
            r, c, m, s, d = map(int, input().split())
            fireball[(r-1, c-1)] = [m, s, [d]]
        
        return fireball
    
    # 이동
    def move(fireball, N):
        new_fireball = {}
        
        for rc, msd in fireball.items():
            r, c = rc
            m, s, d = msd
            for i in d:
                dr, dc = DIR[i]
                nr = (r+dr*s)%N
                nc = (c+dc*s)%N
                if -1 < nr < N and -1 < nc < N:
                    if new_fireball.get((nr, nc)):
                        new_fireball[(nr, nc)].append([m, s, i])
                    else:
                        new_fireball[(nr, nc)] = [[m, s, i]]

        
        return new_fireball
    
    # 파이어볼 합치기
    def merge_fireball(fireball):
        key_lst = fireball.keys()

        new_fireball = {}
        for key in key_lst:
            fire_lst = fireball[key]
            if len(fire_lst) > 1:
                m_lst, s_lst, d_lst = list(map(list, zip(*fire_lst)))
                total_m, total_s = sum(m_lst) // 5, sum(s_lst) // len(s_lst)
                if total_m == 0:
                    continue
                flag_d = fire_lst[0][2] % 2
                for d in d_lst[1:]:
                    if d % 2 != flag_d:
                        total_d = [1, 3, 5, 7]
                        break
                else:
                    total_d = [0, 2, 4, 6]
                
                new_fireball[key] = [total_m, total_s, total_d]
            else:
                m, s, d = fireball[key][0]
                new_fireball[key] = [m, s, [d]]

        return new_fireball

    fireball = INPUT(M)
    for _ in range(K):
        next_fireball = move(fireball, N)
        fireball = merge_fireball(next_fireball)
    
    res = 0
    for m, _, d in fireball.values():
        res += m * len(d)

    print(res)
    
if __name__ == '__main__':
    sol()

최종 코드

• 위에서 언급한 부분들 고려해서 수정
• 시간차는 크게 나지 않음

# 320ms

import sys
input = sys.stdin.readline

# 파이어볼 딕셔너리로 입력받기
def INPUT(M):
    global fireball
    for _ in range(M):
        r, c, m, s, d = map(int, input().split())
        fireball[(r-1, c-1)] = [m, s, [d]]

# 이동
def move(N):
    new_fireball = {}
    global fireball
    
    for rc, msd in fireball.items():
        r, c = rc
        m, s, d = msd
        for i in d:
            dr, dc = DIR[i]
            nr = (r+dr*s)%N
            nc = (c+dc*s)%N
            if new_fireball.get((nr, nc)):
                new_fireball[(nr, nc)].append([m, s, i])
            else:
                new_fireball[(nr, nc)] = [[m, s, i]]

    fireball = new_fireball

def find_next_dir(d_lst):
    flag = d_lst[0] % 2
    for d in d_lst[1:]:
        if d % 2 != flag:
            return 1, 3, 5, 7
    return 0, 2, 4, 6

# 파이어볼 합치기
def merge_fireball():
    global fireball
    key_lst = fireball.keys()

    new_fireball = {}
    for key in key_lst:
        fire_lst = fireball[key]

        if len(fire_lst) <= 1:
            m, s, d = fireball[key][0]
            new_fireball[key] = [m, s, [d]]
            continue
        m_lst, s_lst, d_lst = list(map(list, zip(*fire_lst)))
        total_m, total_s = sum(m_lst) // 5, sum(s_lst) // len(s_lst)
        if total_m == 0:
            continue
        
        new_fireball[key] = [total_m, total_s, find_next_dir(d_lst)]

    fireball = new_fireball

if __name__ == '__main__':
    N, M, K = map(int, input().split())
    DIR = {
        0: (-1, 0),
        1: (-1, 1),
        2: (0, 1),
        3: (1, 1),
        4: (1, 0),
        5: (1, -1),
        6: (0, -1),
        7: (-1, -1)
    }
    fireball = {}
    
    INPUT(M)
    for _ in range(K):
        move(N)
        merge_fireball()

    res = 0
    for m, _, d in fireball.values():
        res += m * len(d)

    print(res)