[백준] 14502번: 연구소 - Python

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

14502번: 연구소

아이디어

• 백트래킹을 통해 벽을 세우는 경우를 만듦
• 그 경우에서 2인 위치를 기준으로 bfs를 진행하여 감염이 된 상태를 만듦
• 이후 안전지대의 개수를 세고 result 값과 비교하여 더 큰 값을 result에 저장.

구현(pypy정답)

import sys
input = sys.stdin.readline

# 백트래킹
def back(wall):
    if wall == 3:
        tmp_li = [mat[:] for mat in li]
        bfs(tmp_li)
        return

    for i in range(n):
        for j in range(m):
            if li[i][j] == 0:
                li[i][j] = 1
                back(wall+1)
                li[i][j] = 0

dy = [0, -1, 0, 1]
dx = [-1, 0, 1, 0]

# bfs진행
from collections import deque
def bfs(tmp_li):
    global result
    q = deque()

    for i in range(n):
        for j in range(m):
            if li[i][j] == 2:
                q.append((i,j))

    while q:
        y, x = q.popleft()
        for i in range(4):
            ny = y + dy[i]
            nx = x + dx[i]
            if 0<= ny <n and 0<=nx<m:
                if tmp_li[ny][nx] == 0:
                    tmp_li[ny][nx] = 2
                    q.append((ny, nx))
    count = check_safe(tmp_li)
    result = max(result, count)

def check_safe(tmp_li):
    cnt = 0
    for i in range(n):
        for j in range(m):
            if tmp_li[i][j] == 0:
                cnt += 1
    return cnt


n, m = list(map(int, input().split()))
li = [list(map(int, input().split())) for _ in range(n)]

result = 0

back(0)

print(result)

pypy로는 통과가 되는데 python3에서는 시간초과가 난다.

원인 분석

• 위 코드는 back 안에서 함수를 돌 때마다 safe지역을 찾는다.
• back함수를 들어가기 전 미리 safe 지역 좌표를 찾고 그걸 통해 for문을 돌면 최소화 할 수 있지 않을까?

구현(정답)

import sys
input = sys.stdin.readline

# 백트래킹
def back(wall, idx):
    if wall == 3:
        tmp_li = [mat[:] for mat in li]
        bfs(tmp_li)
        return
    for x in range(idx,len(safe)):
        li[safe[x][0]][safe[x][1]]=1
        back(wall+1,x+1)
        li[safe[x][0]][safe[x][1]]=0

dy = [0, -1, 0, 1]
dx = [-1, 0, 1, 0]

# bfs진행
from collections import deque
def bfs(tmp_li):
    global result
    q = deque()

    for i in range(n):
        for j in range(m):
            if li[i][j] == 2:
                q.append((i,j))

    while q:
        y, x = q.popleft()
        for i in range(4):
            ny = y + dy[i]
            nx = x + dx[i]
            if 0<= ny <n and 0<=nx<m:
                if tmp_li[ny][nx] == 0:
                    tmp_li[ny][nx] = 2
                    q.append((ny, nx))
    count = check_safe(tmp_li)
    result = max(result, count)

def check_safe(tmp_li):
    cnt = 0
    for i in range(n):
        for j in range(m):
            if tmp_li[i][j] == 0:
                cnt += 1
    return cnt


n, m = list(map(int, input().split()))
li = [list(map(int, input().split())) for _ in range(n)]
safe = []

# 벽을 세울 수 있는 구역 미리 저장
for x in range(n):
    for y in range(m):
        if li[x][y]==0:
            safe.append([x,y])

result = 0

back(0, 0)

print(result)

마무리

• 상수 시간도 백트래킹 안으로 들어가면 곱하기 몇천만이라고 한다.
• 백트래킹의 for문의 크기를 최대한 줄이도록 노력하자.
• for문 크기 줄이는 스킬은 문제를 많이 풀어보고 감을 잡는 것밖에 답이 없는듯함..