[백준 파이썬 14502번 연구소] 구현 최적화

최초 코드

from copy import deepcopy
from collections import deque

dr = [1, 0, -1, 0]
dc = [0, 1, 0, -1]

N, M = map(int, input().split())
graph =[]
for _ in range(N):
    graph.append(list(map(int, input().split())))

zero_ls = []
virus_ls = []
for row in range(N):
    for col in range(M):
        if graph[row][col] == 0:
            zero_ls.append((row, col))
        elif graph[row][col] == 2:
            virus_ls.append((row, col))

def count(visit):
    q = deque(virus_ls)
    while q:
        r, c = q.popleft()
        for i in range(4):
            nr = r + dr[i]
            nc = c + dc[i]
            if nr < 0 or nc < 0 or nr >= N or nc >= M:
                continue
            if visit[nr][nc] == 0:
                visit[nr][nc] = 2
                q.append((nr, nc))

    cnt = 0
    for r in range(N):
        for c in range(M):
            if visit[r][c] == 0:
                cnt += 1
    return cnt

MAX = 0
long = len(zero_ls)
for a in range(long-2):
    for b in range(a+1, long-1):
        for c in range(b + 1, long):
            visited = deepcopy(graph)
            row, col = zero_ls[a]
            visited[row][col] = 1
            row, col = zero_ls[b]
            visited[row][col] = 1
            row, col = zero_ls[c]
            visited[row][col] = 1
            MAX = max(MAX, count(visited))

print(MAX)

무식하게 3중 for문을 통해 조합을 구현

 

그 결과 무려 2680ms!!

그래서 고수들의 코드를 보니

def set_wall(_map, num=1):
    def check_8neighbor(inp_map, n, m):  # 8 neighbor
        score = 0
        for a in range(n - 1, n + 2):
            for b in range(m - 1, m + 2):
                if inp_map[a][b] == 1:
                    score += 1
        if score >= 1:
            return True
        else:
            return False

    N = len(_map)
    M = len(_map[0])
    crnt_map = [x[:] for x in _map]  # deep copy

    maximum = 0
    count = 0
    # set candidate
    candidate = list()
    for n in range(1, N-1):
        for m in range(1, M-1):
            if crnt_map[n][m] == 0 and check_8neighbor(crnt_map, n, m):
                candidate.append((n, m))

    for w in candidate:
        crnt_map[w[0]][w[1]] = 1  # set wall
        if num == 1:
            count = set_wall(crnt_map, num + 1)
            crnt_map[w[0]][w[1]] = 3  # 지나갔던 곳
        if num == 2:
            count = set_wall(crnt_map, num + 1)
            crnt_map[w[0]][w[1]] = 4  # 지나갔던 곳
        if num == 3:
            count = spreadAndCount(crnt_map)  # spread the plague And measure the area of safe zone
            crnt_map[w[0]][w[1]] = 5  # 지나갔던 곳
            if maximum < count:
                maximum = count
            continue

        if maximum < count:
            maximum = count

    return maximum

def isolate(_map):
    M = len(_map[0])
    # padding
    inp_map = [[1 for _ in range(M)]] + _map
    for i, m in enumerate(inp_map):
        inp_map[i] = [1] + m + [1];
    inp_map.append([1 for _ in range(M+2)])

    res = set_wall(inp_map)  # dynamic programming, DFS

    return res

if __name__ == "__main__":
    N, M, _map = get_input()
    result = isolate(_map)
    print(result)

1. graph를 패딩한다
2. 좌표를 그리디로 돌리면서 주변의 벽이 있을 시 candidate를 뽑는다
3. 벽 설치 시 candidate를 대상으로 조합을 짠다

조합이 N의 영향을 많이 받다보니 시간을 많이 줄일 수 있었다.