[2022 KAKAO BLIND RECRUITMENT] 양과 늑대

1트 코드

import sys
sys.setrecursionlimit(10**8)

def DFS(graph, root, info, cnt, visited=None):
    if not graph:
        return []
    if visited is None:
        visited = []
    if info[root] == 0:
        cnt += 1
    elif info[root] == 1:
        cnt -= 1
    if cnt > 0:
        visited.append(root)
        for next_node in graph[root]:
            if next_node not in visited:
                visited = DFS(graph, next_node, info, cnt, visited)
        return visited
    else:
        return [] 


def solution(info, edges):
    answer = 0
    graph = {}

    for edge in edges:
        u, v = edge
        if u not in graph:
            graph[u] = []
        if v not in graph:
            graph[v] = []
        graph[u].append(v)
        graph[v].append(u)

    for root in graph.keys():
        if info[root] == 0:
            result = DFS(graph, root, info, 1)
            answer = max(answer, len(result))

    return answer

2트 코드

def BFS(graph, root, info):
    visited = []
    cnt = 0
    stack = [(root, cnt)]

    while stack:
        node, cnt = stack.pop()

        if info[node] == 0:
            cnt += 1
        elif info[node] == 1:
            cnt -= 1

        if cnt > 0 and node not in visited:
            visited.append(node)
            stack.extend((next_node, cnt) for next_node in graph[node])

    return visited


def solution(info, edges):
    graph = {node: [] for node in range(len(info))}

    for u, v in edges:
        graph[u].append(v)
        graph[v].append(u)

    answer = max(len(BFS(graph, root, info)) for root in graph if info[root] == 0)
    return answer