[백준] 14502번 연구소 - 파이썬(Python)

문제 (링크)

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

14502번: 연구소

 

나의 풀이

import sys
from collections import deque
from itertools import combinations
import copy

input = sys.stdin.readline

n, m = map(int, input().split())
graph = []
virus_locs = []
empty_locs = []
moves = [
    (1, 0),
    (-1, 0),
    (0, 1),
    (0, -1)
]

for i in range(n):
    row = list(map(int, input().split()))
    for j in range(m):
        if row[j] == 2:
            virus_locs.append([i, j])
        
        elif row[j] == 0:
            empty_locs.append([i, j])

    graph.append(row)

def bfs(copied_graph, virus_locs):
    q = deque(virus_locs)

    while q:
        x, y = q.popleft()

        for move in moves:
            nx, ny = x + move[0], y + move[1]

            if 0 <= nx < n and 0 <= ny < m:
                if copied_graph[nx][ny] == 0:
                    copied_graph[nx][ny] = 2
                    q.append([nx, ny])

    return copied_graph

# 빈 공간에서 3개를 뽑는 경우의 수 (3개의 벽을 세우는 경우의 수)
comb = combinations(empty_locs, 3)
ans = -1

for walls in comb:
    x1, y1 = walls[0][0], walls[0][1]
    x2, y2 = walls[1][0], walls[1][1]
    x3, y3 = walls[2][0], walls[2][1]

    # 매번 기존 그래프에서 진행해줘야하므로 copy해서 진행
    copied_graph = copy.deepcopy(graph)

    copied_graph[x1][y1] = 1
    copied_graph[x2][y2] = 1
    copied_graph[x3][y3] = 1

    graph_with_walls = bfs(copied_graph, virus_locs)
    tmp_ans = 0

    for each in graph_with_walls:
        # 각 열에 대해서 0의 갯수 세서 더해줌
        tmp_ans += each.count(0)
    
    ans = max(ans, tmp_ans)

print(ans)

코드 설명 및 풀이법

처음에는 벽을 가장 효과적으로 세울 수 있는 방법에 대해서 오랜 시간 고민했지만, N과 M의 범위가 각각 3이상 8이하라는 점에서 브루트포스 방식을 통해 해결할 수 있는 문제였다.

 

처음 그래프를 입력받을 때 바이러스가 있는 위치와 빈 위치(벽을 세울 수 있는 위치)를 따로 저장해준다.

 

BFS함수에서는 바이러스 위치부터 시작해 BFS를 진행해주고, 진행된 곳에 대해서는 바이러스에 감염됐다는 뜻으로 값을 2로 업데이트 해줬다.

 

앞서 저장해둔 빈 위치에서 3가지 위치를 뽑는 경우의 수를 itertools.combinations함수를 통해서 뽑아준 후, 반복문을 돌며 해당 조합들의 값을 1로 업데이트 해주어 벽을 세워준다.

 

이렇게 벽이 세워진 그래프에 대해서 앞서 정의한 BFS 함수를 돌려준 후 리턴된 그래프에서 0의 갯수를 세어 최대값을 구하도록 작성해준다.