본문 바로가기

Code Review/BaekJoon

[ 백준 21610번 마법사 상어와 비바라기 ] Python 코드 (Simulation)

1. 문제

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

 

21610번: 마법사 상어와 비바라기

마법사 상어는 파이어볼, 토네이도, 파이어스톰, 물복사버그 마법을 할 수 있다. 오늘 새로 배운 마법은 비바라기이다. 비바라기를 시전하면 하늘에 비구름을 만들 수 있다. 오늘은 비바라기

www.acmicpc.net


2. 아이디어

코드 짜기가 어려운 문제는 아닌데, 시간초과가 나버렸다.

두손 모아 기도하며 PyPy로도 돌려보았지만 여전히 시간초과였다.

 

주저리주저리 설명할 것 없이, 내가 개선한 방법만 말하면 다음 3가지다.

  1. 리스트 복사는 하지 않는다. 시간복잡도 O(n)
    리스트의 모든 요소를 복사해야하기 때문에 시간복잡도가 O(n)이다.
  2. 가능하면 리스트 대신 집합에서 in 구문을 쓰자. 시간복잡도 O(n)
    이건 프로그래머스에서 카카오 기출[등산코스 정하기]을 풀 때 안 사실이다.
    할 수 있는 최적화를 다 했는데도 시간초과가 나서 구글링해보니까 list를 set으로 바꿔서 in 구문을 쓰란다.
    반신반의하면서 코드를 수정해서 돌린 결과, 드라마틱한 시간 개선이 있었다.
    나중에 보니까 list에서 in을 쓰면 앞에서부터 순차탐색하기 때문에  시간복잡도가 O(n)인 반면,
    set은 해쉬함수를 통해 각 요소에 해당하는 key값을 계산하기 때문에 딕셔너리처럼 시간복잡도가 O(1)이라고 한다.
    더 자세한 설명은 여기로
  3. 반복문도 최대한 줄이자.
    이건 뭐 당연한다.
    처음 코드를 짤 때는 '구름 이동'과 '비내리기'를 위해 각각 for문을 돌렸다.
    그런데 보니까 한 개의 반복문으로 동시에 수행할 수 있길래 합쳐버렸다.

3. 코드

3.1. 수정 전 코드(시간초과)

import sys
input = sys.stdin.readline

ref = {1: (0, -1), 2: (-1, -1), 3: (-1, 0), 4:(-1, 1), 5:(0, 1), 6:(1, 1), 7:(1, 0), 8:(1, -1)}

N, M = map(int, input().split())
arr = [list(map(int, input().split())) for _ in range(N)]
d = [0] * M
s = [0] * M
for i in range(M):
    d[i], s[i] = map(int, input().split())

# 구름 생성
clouds = [[N-1, 0], [N-1, 1], [N-2, 0], [N-2, 1]]
for m in range(M):
    # 이동
    for cloud in clouds:
        cloud[0] = (cloud[0] + s[m] * ref[d[m]][0]) % N
        cloud[1] = (cloud[1] + s[m] * ref[d[m]][1]) % N

    # 비 내리기
    for ci, cj in clouds:
        arr[ci][cj] += 1

    # 물복사 버그
    for ci, cj in clouds:
        for di, dj in ((-1, -1), (-1, 1), (1, -1), (1, 1)):
            ni = ci + di
            nj = cj + dj
            if 0 <= ni < N and 0 <= nj < N and arr[ni][nj]:
                arr[ci][cj] += 1

    # 구름 생성
    next_clouds = []
    for i in range(N):
        for j in range(N):
            if arr[i][j] >= 2 and [i, j] not in clouds:
                arr[i][j] -= 2
                next_clouds.append([i, j])

    clouds = next_clouds

result = 0
for row in arr:
    result += sum(row)
print(result)

3.2. 수정 후 코드

import sys
input = sys.stdin.readline

ref = {1: (0, -1), 2: (-1, -1), 3: (-1, 0), 4:(-1, 1), 5:(0, 1), 6:(1, 1), 7:(1, 0), 8:(1, -1)}

N, M = map(int, input().split())
arr = [list(map(int, input().split())) for _ in range(N)]
d = [0] * M
s = [0] * M
for i in range(M):
    d[i], s[i] = map(int, input().split())

# 구름 생성
clouds = [(N-1, 0), (N-1, 1), (N-2, 0), (N-2, 1)]

for m in range(M):
    # 이동 후 비내리기
    moved_clouds = set()	# 집합으로 생성
    for i, j in clouds:
        mi = (i + s[m] * ref[d[m]][0]) % N
        mj = (j + s[m] * ref[d[m]][1]) % N
        arr[mi][mj] += 1	# 이동과 비내리기를 동시에 수행
        moved_clouds.add((mi, mj))

    # 물복사 버그
    for ci, cj in moved_clouds:
        for di, dj in ((-1, -1), (-1, 1), (1, -1), (1, 1)):
            ni = ci + di
            nj = cj + dj
            if 0 <= ni < N and 0 <= nj < N and arr[ni][nj]:
                arr[ci][cj] += 1

    # 구름 생성
    clouds = []			# 나중에 복사하지 않고, 그냥 처음부터 clouds에 append
    for i in range(N):
        for j in range(N):
            if (i, j) not in moved_clouds and arr[i][j] >= 2:
                arr[i][j] -= 2
                clouds.append((i, j))

result = 0
for row in arr:
    result += sum(row)
print(result)