PS/Python

[백준] 20056 - 마법사 상어와 파이어볼(파이썬, Python)

w00se 2021. 4. 29. 19:08

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

www.acmicpc.net/problem/20056

 

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

첫째 줄에 N, M, K가 주어진다. 둘째 줄부터 M개의 줄에 파이어볼의 정보가 한 줄에 하나씩 주어진다. 파이어볼의 정보는 다섯 정수 ri, ci, mi, si, di로 이루어져 있다. 서로 다른 두 파이어볼의 위치

www.acmicpc.net

각 파이어 볼은 (행의 위치, 열을 위치, 질량, 속력, 방향)의 정보를 가지고 있으며 각 이동 회차마다 주어진 방향으로 속력만큼 이동을 합니다.

이 문제는 K번 이동 후 남아있는 파이어볼의 질량의 합을 구하는 문제입니다.

 

특징 1

IndexError: list index out of range(인덱스 에러)가 발생하지 않도록 파이어볼의 위치를 계산해야 한다.

 

- 격자의 행과 열은 1번부터 N번까지 번호가 매겨져 있고, 1번 행은 N번과 연결되어 있고, 1번 열은 N번 열과 연결되어 있다.

위의 조건에 의해 파이어볼은 항상 주어진 격자에 존재하며 아래의 예시와 같은 경우가 가능하게 됩니다.

 

예시: 빨간색 점 선 원의 위치에서 5번 방향으로 속력 3만큼 이동시 빨간색 원 위치로 이동

이를 위해 파이어볼의 위치 계산에 나머지 연산을 이용했으며 계산 예시는 아래와 같습니다.

N = 6

#기존 파이어 볼의 정보
r, c, m, s, d = (4, 1, 1, 3, 5)


#이동 후 파이어 볼의 위치 (1, 4)
nr = (r - s) % N
nc = (c - s) % N

 

 

전체 코드는 아래와 같습니다.

 

전체 코드

import sys
import math
import copy

def calculatePos(r, c, s, d):
    nr, nc = 0, 0

    if d == 0:
        nr = (r - s) % N
        nc = c

    elif d == 1:
        nr = (r - s) % N
        nc = (c + s) % N
        
    elif d == 2:
        nr = r
        nc = (c + s) % N
        
    elif d == 3:
        nr = (r + s) % N
        nc = (c + s) % N
        
    elif d == 4:
        nr = (r + s) % N
        nc = c

    elif d == 5:
        nr = (r + s) % N
        nc = (c - s) % N

    elif d == 6:
        nr = r
        nc = (c - s) % N

    else:
        nr = (r - s) % N
        nc = (c - s) % N
        
    return (nr, nc)

N, M, K = map(int, sys.stdin.readline().strip().split())

countMat = [[0] * N for _ in range(N)]
massMat = [[0] * N for _ in range(N)]
speedMat = [[0] * N for _ in range(N)]
directionMat = [[0] * N for _ in range(N)]
multiDirectionMat = [[0] * N for _ in range(N)]

balls = []
answer = 0

for _ in range(M):
    r, c, m, s, d = map(int, sys.stdin.readline().strip().split())
    balls.append((r-1, c-1, m, s, d))

for _ in range(K):
    answer = 0
    newBalls = []
    newPos = []

    for ball in balls:
        r, c, m, s, d = ball

        nr, nc = calculatePos(r, c, s, d)

        massMat[nr][nc] += m
        speedMat[nr][nc] += s
        directionMat[nr][nc] += d

        if countMat[nr][nc] == 0: 
            countMat[nr][nc] += 1
            multiDirectionMat[nr][nc] = d % 2
        
        else:
            countMat[nr][nc] += 1
            
            temp = multiDirectionMat[nr][nc]

            if temp != d % 2:
                multiDirectionMat[nr][nc] = -1

        newPos.append((nr, nc))

    newPos = set(newPos)
    
    for pos in newPos:
        r, c = pos

        if countMat[r][c] == 1:
            m = massMat[r][c]
            s = speedMat[r][c]
            d = directionMat[r][c]

            if m > 0:
                newBalls.append((r, c, m, s, d))
                answer += m

        else:
            m = math.floor(massMat[r][c] / 5)
            s = math.floor(speedMat[r][c] / countMat[r][c])
            if m > 0:
                answer += (m*4)
                if multiDirectionMat[r][c] != -1:
                    for i in range(0, 7, 2):
                        newBalls.append((r, c, m, s, i))
                else:
                    for i in range(1, 8, 2):
                        newBalls.append((r, c, m, s, i))

        countMat[r][c] = 0
        massMat[r][c] = 0
        speedMat[r][c] = 0
        directionMat[r][c] = 0
        multiDirectionMat[r][c] = 0
        
    balls = copy.copy(newBalls)
    
print(answer)