도넛과 막대 그래프 (프로그래머스, 2024 카카오 겨울 인턴십, Python)

풀이

(1) 입력

    MAX_LENGTH = 1_000_001

    go = [[] for _ in range(MAX_LENGTH)]
    take = [[] for _ in range(MAX_LENGTH)]
    for a, b in edges:
        go[a].append(b)
        take[b].append(a)

그래프 모양에 따라 노드에서 나가는 경로와 들어오는 경로의 수가 특정한 규칙을 따르고 있기 때문에 나가는 경로를 나타내는 go 리스트와 take 리스트를 따로 선언합니다. 노드의 범위는 1 <= a, b <= 1,000,000 조건만 있을 뿐, 특정한 조건은 없기 때문에 리스트의 크기를 최댓값인 1_000_001으로 설정합니다.

(2) 시작점 찾기

    start = -1
    for node in range(1_000_002):
        if len(take[node])==0 and len(go[node])>=2:
            start = node
            break

시작점에서는 나가는 경로만 있습니다. 또한 그래프 수의 합이 2 이상이라는 조건 때문에 시작점에서 나가는 경로는 2개 이상임이 보장됩니다. 따라서 들어오는 경로는 하나도 없으면서 나가는 경로는 2개 이상인 노드가 시작점입니다. 다른 노드들의 특성을 볼 때 이런 조건을 만족하는 노드는 시작점 하나로 유일합니다.

(3) BFS 함수에 필요한 리스트 선언

    ans = [start, 0, 0, 0] # 시작, 도넛, 막대, 8자
    visited = [False] * (MAX_LENGTH)

BFS 함수를 정의하기 전에 필요한 리스트를 선언합니다. return 값이 될 ans 리스트와 방문 리스트인 visited를 정의합니다.

(4) BFS 함수 정의

    def bfs(i):
        q = deque([i])
        visited[i] = True

        while q:
            node = q.popleft()
            if not go[node]: # 막대
                ans[2] += 1
                return
            if len(go[node])==2 and len(take[node])==2: # 8자
                ans[3] += 1
                return
            for next_node in go[node]:
                if not visited[next_node]:
                    visited[next_node] = True
                    q.append(next_node) 

        ans[1] += 1

노드 i가 속한 그래프의 모든 노드를 BFS로 방문합니다. 이때 그래프 모양에 따라 다음과 같은 규칙을 갖습니다.

막대 모양 그래프

            if not go[node]: # 막대
                ans[2] += 1
                return

막대 모양 그래프는 마지막 노드에서 나가는 경로가 없습니다. 그래서 이와 같은 조건을 만족하는 노드를 만나면 막대 모양 그래프의 수를 늘리고 return합니다.

8자 모양 그래프

            if len(go[node])==2 and len(take[node])==2: # 8자
                ans[3] += 1
                return

8자 모양 그래프는 중간에 나가는 경로와 들어오는 경로가 모두 2인 노드가 존재합니다. 그래서 이와 같은 조건을 만족하는 노드를 만나면 막대 모양 그래프의 수를 늘리고 return합니다.

도넛 모양 그래프

        ans[1] += 1

막대 모양 그래프와 8자 모양 그래프가 아니라면 도넛 모양 그래프입니다. 그래프의 모든 노드를 BFS로 순회했는데도 아직 return이 되지 않았다면 도넛 모양 그래프로 판단하고 함수를 종료합니다.

(5) BFS 함수 실행

    for node in go[start]:
        take[node].remove(start)
        bfs(node)

각 그래프는 서로 분리되어 있으며 start 노드가 각 그래프를 가리키고 있습니다. 따라서 start 노드가 가리키는 노드에 대해서 BFS 함수를 실행합니다. 단, BFS 함수를 실행하기 전에 start에서 출발한 경로는 삭제합니다.

소스코드

from collections import deque

def solution(edges):
    MAX_LENGTH = 1_000_001

    go = [[] for _ in range(MAX_LENGTH)]
    take = [[] for _ in range(MAX_LENGTH)]
    for a, b in edges:
        go[a].append(b)
        take[b].append(a)

    start = -1
    for node in range(1_000_002):
        if len(take[node])==0 and len(go[node])>=2:
            start = node
            break

    ans = [start, 0, 0, 0] # 시작, 도넛, 막대, 8자
    visited = [False] * (MAX_LENGTH)

    def bfs(i):
        q = deque([i])
        visited[i] = True

        while q:
            node = q.popleft()
            if not go[node]: # 막대
                ans[2] += 1
                return
            if len(go[node])==2 and len(take[node])==2: # 8자
                ans[3] += 1
                return
            for next_node in go[node]:
                if not visited[next_node]:
                    visited[next_node] = True
                    q.append(next_node) 

        ans[1] += 1

    for node in go[start]:
        take[node].remove(start)
        bfs(node)

    return ans