프로그래머스 LV3 - 양과 늑대

문제 링크

내 풀이

# 양 > 늑대
# 방문했는지 안했는지 기록, 방문했으면 양,늑대 수 안더함
# 자신의 자식노드, 부모노드 중에 이동해도 양이 늑대수보다 많으면 이동가능

def is_valid_down(idx, visited, lamb, wolf, info): # 아래로는 안 방문한곳만 가기
    if not visited[idx]:
        if info[idx] == 1:
            if lamb > wolf + 1:
                wolf += 1
                visited[idx] = True
                return True
            else:
                return False
        else:
            lamb += 1
            visited[idx] = True
            return True
    return False

def solution(info, edges):
    lamb = 1 # 양, 늑대 초기 값
    wolf = 0
    parent_dict = {0: 0}
    child_dict = {}
    visited = [False for _ in range(len(info))]
    visited[0] = True
    for edge in edges:
        parent, child = edge
        parent_dict[child] = parent  # 노드 별 부모노드 정보
        if parent not in child_dict: # 노드 별 자식노드 정보
            child_dict[parent] = [child] 
        else:
            child_dict[parent].append(child)
    for i in range(len(info)):
        if i not in child_dict.keys():
            child_dict[i] = None
            
    node = 0

    while True:
        if node == 0 and visited[child_dict[0][0]] == True:
            return lamb
        if child_dict[node] != None or visited[node] != True:
            for child in child_dict[node]:
                if is_valid_down(child, visited, lamb, wolf, info):
                    node = child
                    break
        else:
            node = parent_dict[node] 
    

다른 풀이

from collections import deque

def solution(info, edges):
    def build_tree(info, edges):
        tree = [[] for _ in range(len(info))]
        for edge in edges:
            tree[edge[0]].append(edge[1])
        return tree
    
    tree = build_tree(info, edges)
    max_sheep = 0
    
    q = deque([(0, 1, 0, set())])
    
    while q:
        current, sheep_count, wolf_count, visited = q.popleft()
        max_sheep = max(max_sheep, sheep_count)
        visited.update(tree[current])
        for next_node in visited:
            if info[next_node]:
                if sheep_count != wolf_count + 1:
                    q.append(
                        (next_node, sheep_count, wolf_count + 1, visited - {next_node})
                    )
            else:
                q.append(
                    (next_node, sheep_count + 1, wolf_count, visited - {next_node})
                )
    return max_sheep