삼성 SW역량테스트 기출 / 2019 상반기 오후 2번 문제 바이러스 백신 - BFS, 조합 / Python 파이썬

 

*문제 출처는 삼성전자, 코드트리에 있습니다.

 

삼멘 5일차

문제 제목: 바이러스 백신

문제 사이트: https://www.codetree.ai/training-field/frequent-problems/problems/vaccine-for-virus/description?page=3&pageSize=20&statuses=Ready%2CIn+Progress

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

---

 

 

선택하지 않은 병원에 대해서도 그 뒤에 있는 공간들에 바이러스가 전염되는지 확인해야함

계산 로직 요약

1. 병원과 바이러스 위치 파악:
   • 도시 맵에서 병원(2)과 바이러스(0)의 위치를 저장한다.
2. 병원 조합 선택:
   • 병원들 중 M개의 병원을 선택하는 조합을 찾고, 각 조합에 대해 BFS를 실행한다.
3. BFS로 백신 퍼뜨리기:
   • 선택된 병원들에서 BFS로 백신을 퍼뜨리고, 모든 바이러스를 제거하는 데 걸리는 시간을 계산한다.
4. 최소 시간 계산:
   • 모든 조합을 탐색한 후, 최소 시간을 반환한다. 모든 바이러스를 처리하지 못하면 -1을 출력한다.

 

 

나의 풀이

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

maps = [list(map(int,input().split())) for _ in range(n)]
hospital = [] # 2
visited = [[False] * n for _ in range(n)]

for i in range(n):
    for j in range(n):
        if maps[i][j] == 2:
            hospital.append([i,j])
            maps[i][j] = -2
        elif maps[i][j] == 1:
            visited[i][j] == True
            maps[i][j] = -1
        
            
# m 개의 병원을 고르는 조합
times = 1e9
def combinations(length,arr,c):
    global times
    if len(arr) == length:
        visited_ = [[visited[i][j] for j in range(n)] for i in range(n)]
        maps_ = [[maps[i][j] for j in range(n)] for i in range(n)]
        queue = deque([])
        for i in arr:
            queue.append([i[0],i[1],0])
            visited_[i[0]][i[1]] = True
            maps_[i[0]][i[1]] = 1
        continue_ = False
        for i in range(n):
            for j in range(n):
                if maps[i][j] == 0:
                    continue_ = True
        if not continue_:
            times = 0
            return
        while queue:
            q = queue.popleft()
            qx,qy,cnt = q[0],q[1],q[2]
            direct = [(1,0),(-1,0),(0,1),(0,-1)]
            for i in range(4):
                xx = qx + direct[i][0]
                yy = qy + direct[i][1]
                if 0 <= xx < n and 0 <= yy < n and visited_[xx][yy] == False:
                    if maps_[xx][yy] == 0:
                        maps_[xx][yy] = cnt+1
                        queue.append([xx,yy,cnt+1])
                        visited_[xx][yy] = True
                    elif maps_[xx][yy] == -2:
                        queue.append([xx,yy,cnt+1])
                        visited_[xx][yy] = True
            # for i in maps_:
            #     print(i)
            # print()
        local_times = -1e9
        for i in range(n):
            for j in range(n):
                if maps_[i][j] == 0 and [i,j] not in hospital:
                    return
                if local_times < maps_[i][j]:
                    local_times = maps_[i][j]

        times = min(times,local_times)
        return 
    for i in range(c,len(hospital)):
        if hospital[i] not in arr:
            combinations(length,arr +[hospital[i]],c+1)

combinations(m,[],0)

if times == 1e9 or times == -1e9:
    print(-1)
else:
    print(times)

 

기존에 병원이 있는 공간의 뒤 공간들을 체크하는 과정에서 코드가 난잡해졌다.

그래서 GPT한테 코드 최적화를 부탁함

 

 

모범답안(GPT)

from collections import deque

n, m = map(int, input().split())

maps = [list(map(int, input().split())) for _ in range(n)]
hospital = []  # 병원 위치 저장
virus_count = 0  # 바이러스가 있는 공간 카운트

# 병원 위치와 바이러스 개수 기록
for i in range(n):
    for j in range(n):
        if maps[i][j] == 2:
            hospital.append([i, j])
        elif maps[i][j] == 0:
            virus_count += 1  # 바이러스가 있는 공간의 수

# 방향 이동 설정 (상, 하, 좌, 우)
directions = [(1, 0), (-1, 0), (0, 1), (0, -1)]

# 최소 시간을 구하기 위한 초기값
times = float('inf')

# 병원의 M개 조합을 고르는 함수
def combinations(length, arr, c):
    global times
    if len(arr) == length:
        # 선택된 병원에서 BFS 수행
        bfs(arr)
        return
    for i in range(c, len(hospital)):
        combinations(length, arr + [hospital[i]], i + 1)

# BFS 함수
def bfs(selected_hospitals):
    global times
    queue = deque()
    visited = [[-1] * n for _ in range(n)]  # 방문 기록 및 시간 계산용 리스트
    for hx, hy in selected_hospitals:
        queue.append((hx, hy, 0))  # 병원 위치와 시간 0을 큐에 삽입
        visited[hx][hy] = 0  # 병원 위치는 방문 완료로 표시
    
    total_time = 0  # 걸린 시간
    spread_virus = 0  # 퍼뜨린 바이러스 수

    # BFS 진행
    while queue:
        x, y, time = queue.popleft()

        for dx, dy in directions:
            nx, ny = x + dx, y + dy

            if 0 <= nx < n and 0 <= ny < n and visited[nx][ny] == -1 and maps[nx][ny] != 1:
                visited[nx][ny] = time + 1
                queue.append((nx, ny, time + 1))

                # 바이러스가 있는 공간인 경우
                if maps[nx][ny] == 0:
                    spread_virus += 1
                    total_time = time + 1

    # 바이러스를 모두 퍼뜨렸는지 확인
    if spread_virus == virus_count:
        times = min(times, total_time)

# 병원의 M개 조합 중 최소 시간 계산
combinations(m, [], 0)

# 결과 출력
if times == float('inf'):
    print(-1)
else:
    print(times)

코드 개선 포인트 요약

1. 리스트 복사 제거:
   • visited_, maps_ 리스트를 매번 복사하는 방식은 메모리 낭비가 된다. 이를 줄이기 위해 각 조합마다 BFS 내에서 방문과 시간을 관리하도록 개선한다.
2. 불필요한 플래그 제거:
   • continue_ 플래그는 필요 없다. BFS 내에서 바이러스 퍼뜨린 횟수를 직접 카운팅하면 이 문제를 해결할 수 있다.
3. 최종 시간 계산 방식 개선:
   • BFS 내에서 모든 바이러스를 처리했는지 확인한 후, 불필요한 BFS 반복을 피하고 바로 최소 시간을 기록하도록 한다.

개선 전 코드 (부분)

visited_ = [[visited[i][j] for j in range(n)] for i in range(n)]  # 리스트 복사
maps_ = [[maps[i][j] for j in range(n)] for i in range(n)]
queue = deque([])  
for i in arr:
    queue.append([i[0], i[1], 0])
    visited_[i[0]][i[1]] = True
    maps_[i[0]][i[1]] = 1

개선 후 코드 (부분)

queue = deque([])
visited = [[-1] * n for _ in range(n)]  # 매번 BFS 내에서 visited 초기화
for hx, hy in arr:
    queue.append((hx, hy, 0))  # 병원에서 시작하는 BFS
    visited[hx][hy] = 0  # 병원 위치는 시간 0으로 설정

차이점:

• visited_와 maps_ 리스트 복사를 제거하고, 매번 BFS를 실행할 때 새로 초기화된 visited 리스트만 사용한다.

 

 

 

 

아래: 나, 위: GPT

확실히 시간복잡도 측면에서 크게 개선됐음