[WEEK03] 백준 21606 아침 산책

문제

아침 산책을 즐기는 서현이는 서울과학고에 입학해서도 아침 산책을 즐기려고 합니다. 서현이는 산책을 위해 서울과학고의 지리를 분석했고, 그 결과 서울과학고를
$N$개의 장소를
$N-1$개의 길이 잇는 트리 형태로 단순화시킬 수 있었습니다. 트리 구조이므로, 모든 장소를 몇 개의 길을 통해 오고갈 수 있습니다.

아침 산책은 시작점과 도착점을 정하고, 시작점에서 도착점까지 트리 위의 단순 경로(같은 점을 여러 번 지나지 않는 경로)를 따라 걷게 됩니다. 트리 위의 두 점 사이의 경로는 유일하므로 시작점과 도착점이 정해지면 경로는 유일하게 결정됩니다.

$N$개 장소 중에 일부 장소는 실내이며, 나머지 장소는 실외입니다. 서현이는 산책을 시작하기 전부터 운동을 하는 것을 원치 않기 때문에, 산책의 시작점과 끝점은 모두 실내여야 합니다. 또한, 산책 도중에 실내 장소를 만나면 산책을 그만두고자 하는 욕구가 생기기 때문에, 산책 경로 위에 시작점과 끝점을 제외하고 실내 장소가 있으면 안 됩니다.

서현이는 매일 다른 산책 경로를 걷고자 합니다. 서로 다른 산책 경로가 몇 가지 있는지 구해 봅시다.

입력

첫 줄에는 정점의 수 $N$이 주어집니다.

둘째 줄에는 1과 0으로 이루어진 길이 $N$의 문자열 $A$가 주어집니다. $i$번째 문자 $A_i$가 1일 경우 $i$번 장소는 실내이며, 0인 경우 $i$번 장소는 실외입니다.

셋째 줄부터 $N+1$번 줄까지는 $i+2$번 줄에 트리의 각 간선을 나타내는 두 정수 $u_i$, $v_i$가 주어집니다. 이는 $i$번째 간선이 $u_i$번 정점과 $v_i$번 정점을 연결한다는 의미입니다.

출력

가능한 서로 다른 산책 경로의 수를 출력합니다.

풀이

처음 풀었을 때는 실내 노드일 때 DFS를 진행했으며 다음 노드가 실내라면 카운팅 실외라면 또 탐색을 실시하는 방식으로 하였다.

전체코드 1 (108점 시간초과)

import sys
input = sys.stdin.readline

def dfs(start):
    global count
    visit[start] = 1

    for next in graph[start]:
        if not visit[next]:
            # 실내만나면 카운팅
            if a[next] == '1':
                count += 1
            # 실외면 실내 만날때까지 더 들어가
            else:
                dfs(next)
                visit[next] = 0
        
n = int(input())
a = '0' + input().rstrip()
graph = [[] for _ in range(n + 1)]
visit = [0] * (n + 1)
count = 0

for _ in range(n - 1):
    u, v = map(int, input().split())

    if a[u] == a[v] == 1:
        count += 2
        continue

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

for i in range(1, n + 1):
    # 실내일 때만 시작
    if a[i] == '1':
        dfs(i)
        visit[i] = 0

print(count)

실내끼리 이어진 경우에는 간선정보에 추가하지 않고 바로 2개 카운팅을 해주었다. 방문해제의 경우 총 경우의수를 세는 것이다보니 다른 기준으로 경우를 셀 때를 위해 탐색이 끝났을 때 해제를 해주었다.

시간초과가 났다.

다른 사람의 코드를 참고하였더니 실내 노드 기준으로 탐색을 하지않고 실외 노드를 기준으로 탐색을 하여 실내가 총 3개가 발견된다면 3 * (3 - 1) = 6으로 6가지를 한번에 계산하였다.

그렇다. 실외 기준으로 탐색을 하여 실내를 3개 발견하면 3개중에 2개를 출발노드, 도착노드로 설정하는 3 Permutation 2 -> 3P2를 하면 6이 나온다.

이 방법을 반영한 코드는 다음과 같다.

전체코드 2 (200점) (DFS, BFS)

import sys
from collections import deque
input = sys.stdin.readline

def dfs(start):
    global tmp
    visit[start] = 1

    for next in graph[start]:
        if not visit[next]:
            # 실내만나면 카운팅
            if a[next] == '1':
                tmp += 1
            # 실외면 실내 만날때까지 더 들어가
            else:
                dfs(next)
def bfs(start):
    global tmp
    q = deque()
    visit[start] = 1
    q.append(start)

    while q:
        now = q.popleft()

        for next in graph[now]:
            if not visit[next]:
                if a[next] == '1':
                    tmp += 1
                else:
                    visit[next] = 1
                    q.append(next)
        
n = int(input())
# 인덱스 1부터 쓰려고 '0' 더해준거임
a = '0' + input().rstrip()
graph = [[] for _ in range(n + 1)]
visit = [0] * (n + 1)
count = 0
tmp = 0

for _ in range(n - 1):
    u, v = map(int, input().split())

    # 이어진 두 노드가 실내라면 +2
    # 간선 정보에 안넣어도 됨
    if a[u] == a[v] == '1':
        count += 2
        continue

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

for i in range(1, n + 1):
    # 실외일 때만 시작
    if not visit[i] and a[i] == '0':
        # dfs, bfs 택 1
        # dfs(i)
        bfs(i)
        # 실외 기준으로 실내가 n개 탐색되면
        # nP2 만큼 추가됨
        count += tmp * (tmp - 1)
        tmp = 0

print(count)

DFS, BFS 둘 다 가능하다