Programmers - 양과 늑대

문제 링크

특정 노드를 여러번 방문 할 수 있다는게 중요한 문제인 것 같다.

처음에는 방문 체크하는 배열을 visit[node][sheep][wolf] 로 했는데, 테스트 케이스 한가지를 통과하지 못했다. 어떤게 빠졌는지 확인하다가 visit이 무언가 부족하다는 것을 찾았다.
해당 node에 같은 sheep,wolf 를 끌고 왔어도, 전혀 다른 경로를 통해 왔을 수도 있는데, 이것을 체크하지 않았다.
따라서 방문한 노드 중 가장 큰 노드를 같이 체크하도록 해서 통과할 수 있었다.

def solution(info, edges):
    n = len(info)
    visit = [[[[False for _ in range(n+1)] for __ in range(n+1)]
             for ___ in range(n+1)] for ____ in range(n)]
    g = [[] for _ in range(n)]

    for (x, y) in edges:
        g[x].append(y)
        g[y].append(x)

    answer = [0]

    # 해당 노드에 같은 sheep,wolf로 왔는데, 전혀 다른 경로로 왔을 경우도 있음
    def dfs(node, sheep, wolf, visit_history):
        if visit[node][sheep][wolf][max(visit_history)]:
            return
        if sheep <= wolf:
            return

        visit[node][sheep][wolf][max(visit_history)] = True

        answer[0] = max(answer[0], sheep)
        # print("node", node, (sheep, wolf), visit_history)

        for next in g[node]:
            if not next in visit_history:
                visit_history.add(next)
                if info[next] == 0:
                    dfs(next, sheep+1, wolf, visit_history)
                else:
                    dfs(next, sheep, wolf+1, visit_history)
                visit_history.discard(next)
            else:
                dfs(next, sheep, wolf, visit_history)

    dfs(0, 1, 0, {0})

    return answer[0]