[백준 파이썬] 21610 마법사 상어와 비바리기

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

시간 초과때문에 마지막 새로 생긴 구름의 위치를 생길 수 있는 위치와 사라진 구름의 위치를 반복문 밖으로 빼내 집합으로 구했다.

# 21610 / 마법사 상어와 비바라기
N, M = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(N)]
move = [list(map(int, input().split())) for _ in range(M)]

direction = [ 
    [0, -1], # r, c
    [-1, -1],
    [-1, 0],
    [-1, 1],
    [0, 1],
    [1, 1],
    [1, 0],
    [1, -1],
]
water_copy_dir = [[-1, -1], [-1, 1], [1, -1], [1, 1]]

cloud = [[N-1, 0], [N-1, 1], [N-2, 0], [N-2, 1]]
water_add = [[0 for _ in range(N)] for _ in range(N)]

for d, s in move:
    # 1. 구름 이동
    dd = direction[d-1]
    cloud_moved = []
    for cloud_r, cloud_c in cloud:
        cloud_r = (cloud_r + dd[0] * s) % N
        cloud_c = (cloud_c + dd[1] * s) % N
        # 2. 구름이 있는 곳에 비내림
        A[cloud_r][cloud_c] += 1
        # 물이 증가한 칸 저장
        cloud_moved.append([cloud_r, cloud_c])
    # 4. 증가한 칸 대각선에 물 확인
    for cloud_r, cloud_c in cloud_moved:
        count = 0
        for i in range(4): # 대각에 물이 있는 만큼
            ncr = cloud_r + water_copy_dir[i][0]
            ncc = cloud_c + water_copy_dir[i][1]
            if (0 <= ncr < N and 0 <= ncc < N
                and A[ncr][ncc] > 0):
                count += 1
        A[cloud_r][cloud_c] += count # 물 증가
    # 3. 구름 사라짐
    cloud = []
    for r in range(N):
        for c in range(N):
            if A[r][c] >= 2 and [r, c]:
                cloud.append([r, c])
    # 생성 가능하지만 3.에서 구름이 사라진 위치에는 생성되면 안된다.
    cloud = list(set(map(tuple, cloud)).difference(map(tuple, cloud_moved)))
    for r, c in cloud:
        A[r][c] -= 2
    
print(sum(map(sum, A)))