[알고리즘] 백준 21606 - 아침산책 (파이썬)

실내에서 실내로 가는 경우의 수를 찾는 문제. (중간에 거치는 정점은 무조건 실외여야만 한다.)

무작정 DFS를 쓰는 것이 아니고 조건들을 따져가면서 경우의 수를 찾아야 한다. 처음엔 정점의 개수만큼 반복문을 돌면서 해당 정점에서 DFS를 시작하는 단순한 아이디어로 풀었다가 중간 지점이 실외인 것을 체크하는 부분에서 헤맸고, 결국 다른 사람의 풀이를 보고 이해했다.

아래는 일반적인 DFS 탐색을 하는 60점 짜리 코드

from sys import stdin as s
from sys import setrecursionlimit
setrecursionlimit(10 ** 9)

def DFS(v):
    """ v : 정점의 번호 """
    global cnt

    visited[v] = 1

    for i in graph[v]:
        # 아직 방문하지 않은 노드들에 대해서
        if visited[i] != 1:

            # 해당 노드가 실내라면 도착점으로 지정하고 경로 수 카운트
            if A[i] == 1:
                visited[i] = 1
                cnt += 1
            # 해당 노드가 실외라면 DFS로 재귀 호출
            else:
                DFS(i)

if __name__ == "__main__":
    n = int(s.readline().rstrip()) # 노드의 수
    A = [0] + list(map(int, s.readline().rstrip())) # 노드 n에 대해서 A[n-1]이 1이면 실내, 0이면 실외
    edges = [list(map(int, s.readline().rstrip().split())) for _ in range(n - 1)]
    graph = []

    cnt = 0

    for i in range(n + 1):
        graph.append([])

    for edge in edges:
        u, e = edge

        graph[u].append(e)
        graph[e].append(u)

    for i in range(1, n + 1):
        visited = [0] * (n + 1)

        if A[i] == 1:
            DFS(i)

    print(cnt)

실내를 기준으로 생각하기 보다 실외를 기준으로 생각하고 풀이하면 된다.

실내에서 출발해서 다른 실내로 도착할 수 있는 경우의 수는 N개 중에 하나를 선택하고 나머지 N-1개 중에서 하나를 고르는 것이기 때문에 N(N-1) 로 계산할 수 있다. 이 때 생각해야할 추가적인 경우가 있는데,

 

1. 중간에 실내 노드가 섞여있는 경우 - 실내가 나왔을 때 탐색을 중단하고, 방문하지 않은 노드들을 대상으로 다시 탐색을 진행해준다.

2. 실내에서 실내로 갈 수 있는 경우 - 그래프 형태를 만들어 주는 과정에서 미리 경우의 수를 세어준다.

 

최종 풀이는 아래와 같다.

from sys import stdin as s
from sys import setrecursionlimit
setrecursionlimit(10 ** 9)

def DFS(v):
    """ v : 정점의 번호 """
    indoor = 0

    visited[v] = 1

    for i in graph[v]:
        # 실외일 경우
        if A[i] == 0:
            # 방문하지 않았다면
            if visited[i] == 0:
                visited[i] = 1
                indoor += DFS(i)
        else: 
            indoor += 1

    return indoor

if __name__ == "__main__":
    n = int(s.readline().rstrip()) # 노드의 수
    A = [0] + list(map(int, s.readline().rstrip())) # 노드 n에 대해서 A[n-1]이 1이면 실내, 0이면 실외
    edges = [list(map(int, s.readline().rstrip().split())) for _ in range(n - 1)]
    graph = []

    cnt = 0

    for i in range(n + 1):
        graph.append([])

    for edge in edges:
        u, e = edge

        graph[u].append(e)
        graph[e].append(u)

        # 실내끼리 인접했을 경우 경로를 2개 더해줌.
        if A[u] == 1 and A[e] == 1:
            cnt += 2

    visited = [0] * (n + 1)

    for i in range(1, n + 1):
        indoor = 0

        # 실외를 기준으로 DFS 수행
        if A[i] != 1:
            if visited[i] == 0:
                indoor = DFS(i)

        cnt += indoor * (indoor - 1)

    print(cnt)