[프로그래머스] 양과 늑대(Python)

https://programmers.co.kr/learn/courses/30/lessons/92343

코딩테스트 연습 - 양과 늑대

쉽지 않았던 문제다. 풀면서 생겼던 문제는 현재 상태에서 갈 수 있는 노드를 찾는 과정에서 객체를 같이 참조하면서 무한루프를 돌았다. 그 부분을 어떻게 할지 고민했다.

 

알고리즘은 다음과 같다.

1. 루트 노드부터 시작해서 갈 수 있는 곳이면 감

2. 갈 수 있는 노드라면 갈 수 있는 노드 집합에서 현재 노드를 제거한 객체를 넘긴다

 

갈 수 있는 노드 집합에서 현재 노드를 제거한 객체를 넘기는 방식이 가장 깔끔한 것 같아서 저렇게 풀이했다.

집합 자료형을 사용해서 시간 복잡도도 줄이고(실제로는 update 구문 사용할 때 for loop를 도는 게 더 시간이 적다.) 가독성도 좋게 풀지 않았나.. 나름 뿌듯하다.

 

나의 풀이

from collections import deque
# BFS 풀이
def solution(info, edges):
    graph = [[] for _ in range(len(info))]
    max_sheep = 0
    for edge in edges:
        graph[edge[0]].append(edge[1])
    # 현재위치, 양개수, 늑대개수, 이동 가능한 노드
    q = deque([[0,1,0,set()]])
    while q:
        now, sheepCount, wolfCount, nextNode = q.popleft()
        max_sheep = max(max_sheep,sheepCount)
        nextNode.update(graph[now]) # 이 부분에서 for loop를 사용해서 add를 하는게 시간은 더 짧음
        # set의 Union 연산은 O(len(a) + len(b)) , for loop을 사용하면 O(len(graph[now])) 만큼만
        for next in nextNode:
            if info[next]:
                # 늑대일경우
                if sheepCount != wolfCount + 1:
                    q.append([next,sheepCount,wolfCount+1, nextNode - {next}])
            else:
                q.append([next,sheepCount+1,wolfCount, nextNode - {next}])
    return max_sheep
    
# DFS 풀이
def solution(info, edges):
    graph = [[] for _ in range(len(info))]
    max_sheep = 1
    for edge in edges:
        graph[edge[0]].append(edge[1])
    def dfs(now, sheepCount, wolfCount, nextNode):
        nonlocal max_sheep
        max_sheep = max(max_sheep, sheepCount)
        nextNode.update(graph[now]) 
        for next in nextNode:
            if info[next]:
                #늑대라면
                if sheepCount != wolfCount + 1:
                    dfs(next,sheepCount,wolfCount+1,nextNode-{next})
            else:
                dfs(next,sheepCount+1,wolfCount,nextNode-{next})
    dfs(1,0,0,set())
    return max_sheep