마법사 상어와 파이어

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

20056번: 마법사 상어와 파이어볼

문제요약:

마법사 상어가 특정조건을 따르는 파이어볼에게 이동하라고 명령했을때, 

마법사 상어가 이동을  k번 명령한 후, 남아잉ㅆ는 파이어볼 질량의 합을 구해보자.

 

소요시간:

45분

 

난이도:

골드4

 

제출횟수:

1번

 

코딩:

import sys
input = sys.stdin.readline
from collections import defaultdict

N,M,K = map(int,input().split())

arr = defaultdict(list)

dy = [-1,-1,0,1,1,1,0,-1]
dx = [0,1,1,1,0,-1,-1,-1]

for _ in range(M):
    r,c,m,s,d = map(int,input().split())
    arr[(r-1)*N + (c-1)].append([m,s,d])

for _ in range(K):
    temp = defaultdict(list)
    #이동
    for y in range(N):
        for x in range(N):
            if arr[y*N+x] :
                for m,s,d in arr[y*N+x]:
                    ny = y+dy[d]*s
                    nx = x+dx[d]*s
                    #넘어가면, 원래대로 돌려준다
                    while ny>=N:
                        ny-=N
                    while ny<0:
                        ny+=N
                    while nx>=N:
                        nx-=N
                    while nx<0:
                        nx+=N
                    temp[ny*N + nx].append([m,s,d])
                    
    #겹치는게 있다면, 조건에 맞게 분리시켜주기
    for y in range(N):
        for x in range(N):
            if len(temp[y*N+x])>=2:
                m,s,d= zip(*temp[y*N+x])
                m = sum(m)//5
                if m==0:
                    temp[y*N+x] = []
                    continue
                s = sum(s)//len(temp[y*N+x])
                piv =1
                #모두짝수거나 모두 홀수라면
                if all([i%2 for i in d]) or not any([i%2 for i in d]):
                    piv =0
                #초기화하고
                temp[y*N+x] = []
                for dir in range(0,7,2):
                    temp[y*N+x].append([m,s,dir+piv])
    arr = temp

ans = 0

for y in range(N):
    for x in range(N):
        if arr[y*N+x]:
            if len(arr[y*N+x])==1:
                ans += arr[y*N+x][0][0]
            else:
                m,s,d= zip(*temp[y*N+x])
                ans += sum(m)

print(ans)

이번 문제 역시 구현 문제기 때문에, 알고리즘을 알아야 하진 않았다. 하지만 defaultdict이나 all, any,zip등 함수들을 통해 조금더 간결하게 문제를 구현할 수 있었다. 

시간을 끈 부분이 m,s,d를 따로 추출할때 zip(*temp)를 사용했다.

만약 [[1,2,3],[4,5,6]]이temp라면 zip(*)함수를 사용하면 [(1,4),(2,5),(3,6)] 과 같이 열을 추출할 수 있는데, 각열을 추출하기 위해 for m,s,d in zip(*temp)를 사용했어서 오류가 있었다 이렇게 하면, m에 (1,4) s에 (2,5) d에(3,6)이 아닌 m에 1 s에 4가 할당되고 d에 할당될것이 없어서 오류가 뜨게된다. m,s,d = zip(*temp)와같이 해야 원하는데로 할당할 수 있다.