[Python/파이썬] 백준 21611번 마법사 상어와 블리자드

21611번: 마법사 상어와 블리자드

 

문제

생략

입력

첫째 줄에 N, M이 주어진다. 둘째 줄부터 N개의 줄에는 격자에 들어있는 구슬의 정보가 주어진다. r번째 행의 c번째 정수는 (r, c)에 들어있는 구슬의 번호를 의미한다. 어떤 칸에 구슬이 없으면 0이 주어진다. 상어가 있는 칸도 항상 0이 주어진다.

다음 M개의 줄에는 블리자드 마법의 방향 di와 거리 si가 한 줄에 하나씩 마법을 시전한 순서대로 주어진다.

출력

첫째 줄에 1×(폭발한 1번 구슬의 개수) + 2×(폭발한 2번 구슬의 개수) + 3×(폭발한 3번 구슬의 개수)를 출력한다.

코드

from collections import deque
n, m = map(int, input().split())

ans  = [0] * 3  # 폭발한 i번 구슬 개수

# 구슬 정보 입력
marble = []
for _ in range(n):
    marble.append(list(map(int, input().strip().split())))
    
dx = [0, 1, 0, -1]
dy = [-1, 0, 1, 0]
    
# 빈 칸이 생겨서 구슬이 이동
def move(func):
    marble_list = func()
    x, y = n//2, n//2   # 초기 위치(상어)
    dir = 0
    cnt = 1
    for _ in range(n):
        for _ in range(2):  # 같은 거리로는 2번 이동
            for _ in range(cnt):
                x += dx[dir%4]
                y += dy[dir%4]
                if 0 <= x < n and 0 <= y < n:
                    if marble_list:
                        marble[x][y] = marble_list.popleft()
                    else:
                        marble[x][y] = 0
            dir += 1        # 왼쪽으로 회전
        cnt += 1            # 거리 증가

# 4개 이상 연속하는 구슬 폭발
def explosion():
    x, y = n//2, n//2   # 초기 위치(상어)
    dir = 0
    cnt = 1
    marble_cnt = 1
    prev = -1
    remain_marble = []
    for _ in range(n):
        for _ in range(2):  # 같은 거리로는 2번 이동
            for _ in range(cnt):
                x += dx[dir%4]
                y += dy[dir%4]
                if 0 <= x < n and 0 <= y < n:
                    if marble[x][y] == prev:
                        marble_cnt += 1
                    else:
                        if marble_cnt < 4:
                            if prev > 0:
                                for _ in range(marble_cnt):
                                    remain_marble.append(prev)
                        else:
                            ans[prev-1] += marble_cnt  # 폭발한 구슬
                        prev = marble[x][y]
                        marble_cnt = 1
            dir += 1
        cnt += 1
        
    flag = True
    while flag:
        if not remain_marble:   # 남은 구슬이 없는 경우
            return remain_marble
        cnt = 1
        prev = remain_marble[0]
        tmp = []
        flag = False
        for m in remain_marble[1:]:
            if prev != m:
                if cnt < 4:
                    for _ in range(cnt):
                        tmp.append(prev)
                else:   # 폭발하는 구슬
                    flag = True
                    ans[prev-1] += cnt
                prev = m
                cnt = 1
            else:
                cnt += 1
        # 남은 구슬 처리
        if cnt < 4:
            for _ in range(cnt):
                tmp.append(prev)
        else:   # 폭발하는 구슬
            flag = True
            ans[prev-1] += cnt
        remain_marble = tmp
    return deque(remain_marble)

# 구슬 변화
def change():
    x, y = n//2, n//2   # 초기 위치(상어)
    dir = 0
    cnt = 1
    marble_cnt = 1
    prev = -1
    remain_marble = deque([])
    for _ in range(n):
        for _ in range(2):  # 같은 거리로는 2번 이동
            for _ in range(cnt):
                x += dx[dir%4]
                y += dy[dir%4]
                if 0 <= x < n and 0 <= y < n:
                    if marble[x][y] == prev:
                        marble_cnt += 1
                    else:
                        if prev > 0:
                            remain_marble.append(marble_cnt)
                            remain_marble.append(prev)
                        prev = marble[x][y]
                        marble_cnt = 1
            dir += 1
        cnt += 1
    return remain_marble

# 마법 시전 정보 입력
mx = [-1, 1, 0, 0]
my = [0, 0, -1, 1]
for i in range(m):
    d, s = map(int, input().strip().split())
    x, y = n//2, n//2
    for _ in range(s):
        x += mx[d-1]
        y += my[d-1]
        marble[x][y] = 0    
    move(explosion)
    move(change)
    
print(ans[0] + 2 * ans[1] + 3 * ans[2])

남은 구슬들을 문제에 주어진 칸의 순서대로 1차원 배열로 만들어서 리턴해주는 함수들 사용.

• marble
  :격자의 구슬 정보를 저장하는 2차원 배열
• move()
  : 파라미터로 받은 함수들은 격자에 남은 구슬들을 칸의 순서대로 1차원 배열에 담아 리턴한다. move 함수는 이 배열을 이용하여 marble 배열을 갱신해준다.
• explosion()
  : 폭발 후 남은 구슬 리턴.
  이 때 1차 폭발 후에도 또 구슬이 폭발할 수 있다 생성된 remain_marble 배열의 길이가 0이 되거나 더 이상 폭발이 없을 때까지 remain_marble 배열을 탐색하며 4개 연속된 구슬이 폭발하는지 확인한다.
• change()
  : 구슬 변화 후의 구슬들 배열 리턴.

 

문제 제출 후 64퍼센트 쯤에서 인덱스 에러가 나서 왜 그런지 살펴보았더니 explosion 함수에서 1차 폭발 이후 연속되는 폭발들을 처리해줄 때 remain_marble의 길이가 0이 되는 경우를 고려해주지 않아서 그런 것이었다. 

if not remain_marble:   # 남은 구슬이 없는 경우
	return remain_marble

이걸 추가해줬더니 해결됐다.

 

처음에는 마법 시전 직후의 구슬들의 상태를 리턴하는 get_remain_marble이라는 함수를 만들고 마법 시전 직후 파괴된 구슬을 처리한 marble 2차원 배열에 대해 explosion을 다시 했는데 생각해보니 explosion에서 마법 시전으로 없어져버린 구슬들은 marble에서 0으로 처리되어서 배열에 안 들어오니 굳이 move를 불필요하게 한 번 더 할 필요가 없다 싶어서 get_reamin_marble 부분은 뺐다. 빼도 잘 돌아간다. 아래의 코드가 get_remain_marble 함수가 있던 코드이다.

from collections import deque
n, m = map(int, input().split())

ans  = [0] * 3  # 폭발한 i번 구슬 개수

# 구슬 정보 입력
marble = []
for _ in range(n):
    marble.append(list(map(int, input().strip().split())))
    
dx = [0, 1, 0, -1]
dy = [-1, 0, 1, 0]

# 마법 시전 후 남은 구슬들 칸 순서대로 구하기 => 없어도 됨
def get_remain_marble():
    x, y = n//2, n//2   # 초기 위치(상어)
    dir = 0
    cnt = 1
    remain_marble = deque([])
    for _ in range(n):
        for _ in range(2):  # 같은 거리로는 2번 이동
            for _ in range(cnt):
                x += dx[dir%4]
                y += dy[dir%4]
                if 0 <= x < n and 0 <= y < n:
                    if marble[x][y] > 0:
                        remain_marble.append(marble[x][y])
            dir += 1    # 왼쪽으로 회전 
        cnt += 1    # 이동 거리 증가
    return remain_marble
    
# 빈 칸이 생겨서 구슬이 이동
def move(func):
    marble_list = func()
    x, y = n//2, n//2   # 초기 위치(상어)
    dir = 0
    cnt = 1
    for _ in range(n):
        for _ in range(2):  # 같은 거리로는 2번 이동
            for _ in range(cnt):
                x += dx[dir%4]
                y += dy[dir%4]
                if 0 <= x < n and 0 <= y < n:
                    if marble_list:
                        marble[x][y] = marble_list.popleft()
                    else:
                        marble[x][y] = 0
            dir += 1    # 왼쪽으로 회전
        cnt += 1    # 이동 거리 증가

# 4개 이상 연속하는 구슬 폭발
def explosion():
    x, y = n//2, n//2   # 초기 위치(상어)
    dir = 0
    cnt = 1
    marble_cnt = 1
    prev = -1
    remain_marble = []
    for _ in range(n):
        for _ in range(2):  # 같은 거리로는 2번 이동
            for _ in range(cnt):
                x += dx[dir%4]
                y += dy[dir%4]
                if 0 <= x < n and 0 <= y < n:
                    if marble[x][y] == prev:
                        marble_cnt += 1
                    else:
                        if marble_cnt < 4:
                            if prev > 0:
                                for _ in range(marble_cnt):
                                    remain_marble.append(prev)
                        else:
                            ans[prev-1] += marble_cnt  # 폭발한 구슬
                        prev = marble[x][y]
                        marble_cnt = 1
            dir += 1    # 왼쪽으로 회전
        cnt += 1    # 이동 거리 증가
    
    # 1차 폭발 후 연속되는 폭발들 처리
    flag = True
    while flag: # 폭발이 더이상 없을 때까지 반복
        if not remain_marble:   # 남은 구슬이 없는 경우
            return remain_marble
        cnt = 1
        prev = remain_marble[0]
        tmp = []
        flag = False
        for m in remain_marble[1:]:
            if prev != m:
                if cnt < 4: # 4개 연속X -> 폭발 안 함
                    for _ in range(cnt):
                        tmp.append(prev)
                else:       # 4개 연속 O -> 폭발
                    flag = True
                    ans[prev-1] += cnt
                prev = m
                cnt = 1
            else:
                cnt += 1
        # 마지막 남은 구슬 처리
        if cnt < 4: # 4개 연속X -> 폭발 안 함
            for _ in range(cnt):
                tmp.append(prev)
        else:       # 4개 연속 O -> 폭발
            flag = True
            ans[prev-1] += cnt
        remain_marble = tmp
    return deque(remain_marble)

# 구슬 변화
def change():
    x, y = n//2, n//2   # 초기 위치(상어)
    dir = 0
    cnt = 1
    marble_cnt = 1
    prev = -1
    remain_marble = deque([])
    for _ in range(n):
        for _ in range(2):  # 같은 거리로는 2번 이동
            for _ in range(cnt):
                x += dx[dir%4]
                y += dy[dir%4]
                if 0 <= x < n and 0 <= y < n:
                    if marble[x][y] == prev:
                        marble_cnt += 1
                    else:
                        if prev > 0:
                            remain_marble.append(marble_cnt)
                            remain_marble.append(prev)
                        prev = marble[x][y]
                        marble_cnt = 1
            dir += 1    # 왼쪽으로 회전
        cnt += 1    # 이동 거리 증가
    return remain_marble

# 마법 시전 정보 입력
mx = [-1, 1, 0, 0]
my = [0, 0, -1, 1]
for i in range(m):
    d, s = map(int, input().strip().split())
    x, y = n//2, n//2
    for _ in range(s):  # 마법 시전
        x += mx[d-1]
        y += my[d-1]
        marble[x][y] = 0    
    move(get_remain_marble) # 블리자드 마법 시전 직후 이동
    move(explosion)         # 4개 이상 연속하는 구슬 폭발  
    move(change)            # 구슬 변화
    # print("----%i번째 블리자드----"%(i+1))
    # for j in range(n):
    #     print(*marble[j])
    
print(ans[0] + 2 * ans[1] + 3 * ans[2])