[Python] 백준 14502: 연구소 (조합, bfs 최적화)

문제 링크

14502번: 연구소

문제 요약

연구소에 바이러스가 유출되었다. 바이러스는 벽을 뚫지 못하며 가로 또는 세로 한 칸씩 이동할 수 있다. 벽을 3개 세울 때, 확보할 수 있는 안전 공간의 최대 갯수를 구하고 싶다. 이 때 안전 공간은 벽으로 막혀서 바이러스가 오지 못하는 공간이다.
입력으로 N×M 크기의 연구소가 주어진다. 0은 빈공간, 1은 벽, 2는 바이러스를 의미한다. 

풀이 과정

예시의 입출력에서 안전 공간을 최대로 확보할 수 있는 최적의 벽 위치를 컴퓨터가 어떻게 알 수 있을까 생각해본 후,

1) 모든 벽을 세우는 경우에 대해서

2) 안전 공간의 넓이를 계산하고 이의 최대값을 구한다.

라는 접근으로 풀 수밖에 없다고 생각이 들었다.

 

첫 번째 시도: for문으로 조건에 맞는 경우를 찾아 탐색

우선, 전체 코드는 아래와 같다.

from copy import deepcopy
from collections import deque
N, M=map(int, input().split())
virus_map=[]
for _ in range(N):
    virus_map.append(list(map(int, input().split())))

def safety_zone(virus_map):
    dx=[1, -1, 0, 0]
    dy=[0, 0, 1, -1]
    visited=[[0 for _ in range(M)] for _ in range(N)]

    for i in range(N):
        for j in range(M):
            if virus_map[i][j]==2 and visited[i][j]==0:
                q=deque([(i, j)])
                while q:
                    x, y=q.popleft()
                    visited[x][y]=1
                    for k in range(4):
                        cx, cy=x+dx[k], y+dy[k]
                        if 0<=cx<N and 0<=cy<M:
                            if visited[cx][cy]==0:
                                if virus_map[cx][cy]==0:
                                    virus_map[cx][cy]=2
                                    q.append((cx, cy))
                                elif virus_map[cx][cy]==1:
                                    continue
                                elif virus_map[cx][cy]==2:
                                    q.append((cx, cy))

    safety_zone_size=0
    for i in range(N):
        for j in range(M):
            if virus_map[i][j]==0:
                safety_zone_size+=1
    
    return safety_zone_size

max_safety_zone=0
for i in range(N*M-2):
    x1, y1=i//M, i%M
    if virus_map[x1][y1]>0:
        continue
    for j in range(i+1, N*M-1):
        x2, y2=j//M, j%M
        if virus_map[x2][y2]>0:
            continue
        for k in range(j+1, N*M):
            x3, y3=k//M, k%M
            if virus_map[x3][y3]>0:
                continue
            virus_map_arg=deepcopy(virus_map)
            virus_map_arg[x1][y1]=1
            virus_map_arg[x2][y2]=1
            virus_map_arg[x3][y3]=1
            max_safety_zone=max(max_safety_zone, safety_zone(virus_map_arg))

print(max_safety_zone)

 

위 코드를 두 부분으로 나누어 보면, main부는 아래와 같다.

가능한 벽을 세우는 모든 경우의 수는, 삼중 for문을 돌면서 세 좌표에서 virus_map의 값이 모두 0일 때 max_safety_zone을 업데이트 하는 방식으로 탐색하였다.

max_safety_zone=0
for i in range(N*M-2):
    x1, y1=i//M, i%M
    if virus_map[x1][y1]>0:
        continue
    for j in range(i+1, N*M-1):
        x2, y2=j//M, j%M
        if virus_map[x2][y2]>0:
            continue
        for k in range(j+1, N*M):
            x3, y3=k//M, k%M
            if virus_map[x3][y3]>0:
                continue
            virus_map_arg=deepcopy(virus_map)
            virus_map_arg[x1][y1]=1
            virus_map_arg[x2][y2]=1
            virus_map_arg[x3][y3]=1
            max_safety_zone=max(max_safety_zone, safety_zone(virus_map_arg))

 

각 좌표의 virus_map 값이 0인지 확인하고, 조건에 맞는 세 벽의 좌표를 찾으면 해당 위치에 벽이 세워졌을 때 안전 구역의 넓이를 구하는 safety_zone을 실행하여 max_safety_zone을 업데이트한다.

 

safety_zone은 bfs 알고리즘 방식으로 virus[i][j]==2일 때, 해당 좌표를 queue에 넣고 바이러스가 전파된 후의 map을 구한다. 이 때, visited matrix를 사용하여 중복하여, 무한 탐색되지 않도록 처리하고자 하였다.

(이는 초기 시도로써 특히 virus_map[cx][cy]==2: q.append(cx, cy)는 불필요한 코드이다.)

def safety_zone(virus_map):
    dx=[1, -1, 0, 0]
    dy=[0, 0, 1, -1]
    visited=[[0 for _ in range(M)] for _ in range(N)]

    for i in range(N):
        for j in range(M):
            if virus_map[i][j]==2 and visited[i][j]==0:
                q=deque([(i, j)])
                while q:
                    x, y=q.popleft()
                    visited[x][y]=1
                    for k in range(4):
                        cx, cy=x+dx[k], y+dy[k]
                        if 0<=cx<N and 0<=cy<M:
                            if visited[cx][cy]==0:
                                if virus_map[cx][cy]==0:
                                    virus_map[cx][cy]=2
                                    q.append((cx, cy))
                                elif virus_map[cx][cy]==1:
                                    continue
                                elif virus_map[cx][cy]==2:
                                    q.append((cx, cy))

    safety_zone_size=0
    for i in range(N):
        for j in range(M):
            if virus_map[i][j]==0:
                safety_zone_size+=1
    
    return safety_zone_size

두 번째 시도: bfs 탐색 중 중단 조건 만들기

장렬하게 시간 초과를 맞고, 소요 시간을 줄일 수 있는 방법을 고민하였다.

3번의 시도 끝에 성공한 14502번 문제

 

bfs 탐색 중, max_safety가 될 수 없는 상황일 때 탐색을 중단하는 방식으로 코드를 짜보자는 아이디어로, safety_zone을 구했던 것과 반대로 min_danger_zone을 업데이트하고, 마지막에 원본 virus_map의 safety_zone 사이즈에 (dager_zone 사이즈+벽의 갯수인 3)을 빼주는 방식으로 다시 코드를 작성하였다.

(처음에 벽 개수 빼주는 것을 생각 못해서 답이 조금씩 달라 당황했다. 예시 입출력 정답의 벽 위치 조건일 때, 바이러스 전파 후 virus_map을 출력하고 이를 danger_zone_size 등과 비교하면서 문제를 발견했다.)

min_danger_zone=min(min_danger_zone, danger_zone(virus_map_arg, min_danger_zone))

 

아래는 danger_zone을 구하는 함수의 코드이다. 0=>2로 변경한 경우에 danger_zone을 1씩 더한다.

if danger_zone_size>min_danger_zone이면, 더이상의 탐색이 의미가 없기 때문에 곧바로 return 한다. 

def danger_zone(virus_map, min_danger_zone):
    dx=[1, -1, 0, 0]
    dy=[0, 0, 1, -1]
    visited=[[0 for _ in range(M)] for _ in range(N)]

    danger_zone_size=0
    for i in range(N):
        for j in range(M):
            if virus_map[i][j]==2 and visited[i][j]==0:
                q=deque([(i, j)])
                while q:
                    x, y=q.popleft()
                    visited[x][y]=1
                    for k in range(4):
                        cx, cy=x+dx[k], y+dy[k]
                        if 0<=cx<N and 0<=cy<M:
                            if visited[cx][cy]==0:
                                if virus_map[cx][cy]==0:
                                    virus_map[cx][cy]=2
                                    q.append((cx, cy))
                                    danger_zone_size+=1
                                    if danger_zone_size>min_danger_zone:
                                        return danger_zone_size
                                elif virus_map[cx][cy]==1:
                                    continue
                                elif virus_map[cx][cy]==2:
                                    q.append((cx, cy))

    return danger_zone_size

 

이 코드도 마찬가지로 시간 초과가 떴다.

 

다시 찬찬히 입출력을 살펴보니, 아래의 예제는 바꾼 코드에서도 마찬가지로 비효율적으로 돌아가겠다 생각되었다. 실제로 local에서 실행해봐도 꽤 시간이 많이 걸리는 입출력 예시였다.

소요 시간이 가장 긴 까다로운 입출력 예시

세 번째 시도: 삼중 for문→ 단일 for문 순환, visited matrix 밑 판단 조건문 제거

아이디어가 떠오르지 않아 다른 답 코드를 찾아보았다.

아래의 포스팅에서 많은 도움을 얻었는데, 코드의 구성 자체는 동일했지만 뒤의 최적화 방법을 차용함으로써 시간 초과 문제를 해결할 수 있었다.

[BFS/Combination] 백준#14502: 연구소/Python 풀이

 

첫 번째로, 가능한 모든 벽들의 좌표 위치는 원본 virus_map에서 value가 0인 좌표들 중 3개를 선택하는 조합과 같을 것이다. 기존에는 이를 삼중 for문으로 순환하면서 탐색하였으나, 위 포스팅에서 소개된 itertools의 combinations를 사용하였다. combinations 메소드가 특별히 시간 복잡도가 낮은 것은 아닌데 경험 상 재귀 함수나, 다중 for문처럼 메모리 위치 이동(?) 빈번한 코드가 소요 시간이 대략적으로 예측한 시간보다 길게 걸리는 것 같아서 소요 시간 최적화에 도움을 주지 않았나 추측한다.

 

두 번째로, 초기 queue에 모든 virus의 위치 좌표를 넣어서 이후에 virus_map[i][j]==2일 때, 따로 visited를 확인하거나 queue.append()하는 절차를 생략하였다. 이 때, 위의 참고 코드에서 마음에 들었던 부분이 아래와 같이 list comprehension으로 초기 queue를 간결하게 선언한다는 점이었다.

q=[(i, j) for i in range(N) for j in range(M) if virus_map[i][j]==2]

 

이 두 가지 부분을 수정한 뒤, local에서 실행했을 때에 훨씬 코드의 실행 속도가 빨라짐을 느낄 수 있었고, 백준 채점도 통과하였다.

bfs 관련 문제들을 풀면서 한 가지 느끼는 점은, 기본적인 알고리즘를 그대로 쓰기보다는 각 문제에 따라 최적화 가능한 부분은 최적화를 하면서 조금씩 다르게 코드를 쓰는 것 같다. 그런데, 최적화에 너무 신경을 쓰다보면 실수가 많아지고, 원하는 return 값이 안 나올 때 수정해야하는 부분을 찾기가 어려워져 적당히가 필요한 것 같다. 

 

알고리즘을 적당히 활용하는 것은 결국 쌓은 경험이 많아져야 더 쉬워지지 않을까하고 생각한다.