[코테] 프로그래머스 - 양과 늑대 Python 코드

 

def solution(info, edges):
    info = [1 if x == 0 else -1 for x in info]  # 1: 양, -1: 늑대
    graph = [set() for _ in range(len(info))]   # i번 노드에서 갈 수 있는 노드들
    for [p, c] in edges:
        graph[p].add(c)

    def visit(node, info, graph, sumCnt, sheepCnt, nextNodes):
        if sumCnt + info[node] == 0:    # 양, 늑대 수가 같으면 방문 종료
            return
        
        sumCnt += info[node]
        if info[node] == 1:
            sheepCnt += 1
            nonlocal maxSheepCnt
            maxSheepCnt = max(sheepCnt, maxSheepCnt)    # 최대 양의 수 갱신

        nextNodes.discard(node)         # 방문할 노드 리스트에서 현재 노드 제거
        nextNodes.update(graph[node])   # 방문할 노드 리스트에 갈 수 있는 자식 노드 추가

        for nextNode in nextNodes:
            visit(nextNode, info, graph, sumCnt, sheepCnt, nextNodes.copy())    # copy 주의
    
    maxSheepCnt = 0
    visit(0, info, graph, 0, 0, set())
    return maxSheepCnt
    
    
# https://school.programmers.co.kr/learn/courses/30/lessons/92343