[WEEK 03] 백준 21606번 - 아침 산책 (Python)

백준 21606번 : 아침 산책

21606번: 아침 산책

 

주어진 조건을 만족하는 서로 다른 산책 경로의 개수를 찾는 문제이다. 문제 조건을 요약하면 다음과 같다.

 

 

• 시작점과 끝점은 모두 실내 장소
• 산책 경로 위에 시작점과 끝점을 제외하고 실내 장소가 있으면 안됨

 

 

 

실내 → 실외 →  실내로 이동하는 경우와 실내 → 실내로 이동하는 경우 두 가지를 고려해야 한다.

 

 

 

실내 → 실외 → 실내

 

먼저 하나의 실외 노드에 4개의 실내 노드가 인접해 있다고 가정해보자. 산책 경로의 수는 4개의 실내 노드 중 2개를 순서를 고려하여 뽑는 경우의 수와 동일하다. 따라서 4 × 3 = 12이다. 즉, 실외 노드의 인접한 노드 중 실내 노드의 개수를 세면 산책 경로의 수를 계산할 수 있다.

 

 

실외 노드가 여러 개인 경우도 고려해야 하는데, 이때 산책 경로의 수는 여러 개의 실외 노드를 묶어서 하나의 실외 노드로 가정하여 생각했을 때와 동일하다. 따라서 재귀 호출을 통해 인접한 실외 노드로 이동하여 실내 노드 개수를 세는 행위를 반복하면 된다.

 

 

 

실내 → 실내

 

노드가 실내인지 실외인지 구분할 수 있고, 간선이 주어진다. 인접한 두 노드가 모두 실내 노드이면 산책 경로이므로, 인접 리스트를 형성할 때 해당 산책 경로의 개수를 셀 수 있다. 간선으로 이어진 노드를 인접 리스트에 추가하면서 해당 과정을 수행할 수 있다.

 

 

import sys

sys.setrecursionlimit(10**6)

sys.stdin = open('input.txt', 'r')

input = sys.stdin.readline


def dfs(k, count) :
    visited[k] = 1

    for i in tree[k] :
        if A[i] == 1 :
            count += 1

        elif A[i] == 0 and not visited[i] :
            count = dfs(i, count)           # 재귀 함수의 return 값을 변수에 저장
                                            # if 문에서 update된 count 변수가 인수로 들어감
    return count


N = int(input())

A = [0] + list(map(int, input().strip()))

tree = [[] for _ in range(N + 1)]


case1 = 0

for _ in range(N - 1) :
    u, v = map(int, input().split())
    tree[u].append(v)
    tree[v].append(u)

    if A[u] == 1 and A[v] == 1:
        case1 += 2

visited = [0] * (N + 1)


case2 = 0

for j in range(1, N + 1) :
    if A[j] == 0 and not visited[j] :
        res = dfs(j, 0)                     # 다른 실외 노드의 인접한 실내 노드를 찾는 것이므로 0부터 다시 시작
        case2 += res * (res - 1)            # 인접한 실내 노드를 모두 찾았을 때마다 계산해주어야 함

path = case1 + case2

print(path)