[Python] 백준 21606번 아침 산책

문제

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

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

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

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

입력

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

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

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

출력

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

제한

•  2≤N≤2×105$2 \le N \le 2 \times 10^5$ 
•  1≤ui,vi≤N$1 \le u_i , v_i \le N$ 
•  ui≠vi$u_i \neq v_i$ 
• 입력으로 주어지는 구조는 올바른 트리를 형성함이 보장됩니다.

서브태스크

번호배점제한

1	3	Ai=0$A_i = 0$
2	35	1≤i<N$1 \le i < N$인 모든 정수 i$i$에 대해, 트리의 i$i$번 정점과 i+1$i+1$번 정점 사이에 간선이 존재합니다.
3	10	Ai=1$A_i = 1$인 i$i$는 2개 이하입니다.
4	60	N≤1000$N \le 1000$
5	92	추가 제한 조건이 없습니다.

 

21606번: 아침 산책

---

나의 풀이

이 문제는 처음에 for문을 실내 기준으로 했더니 시간 초과났다. 

import sys
from collections import defaultdict
sys.setrecursionlimit(1000000000)
n = int(sys.stdin.readline().rstrip())
a = list(map(int, list(sys.stdin.readline().strip())))
check = [1]
for i in a:
    check.append(i)

graph = defaultdict(list)

for i in range(n-1):
    s, e = map(int, sys.stdin.readline().split())
    graph[s].append(e)
    graph[e].append(s)


# 실내마다 root를 넣어서 DFS를 넣어주며 경로를 센다. 
cnt = 0
def DFS(start):
    global cnt, visited
    visited[start] = True
    for i in graph[start]:
        if check[i] ==1 and not visited[i]:
            visited[i] = True
            cnt += 1
        elif check[i] ==0 and not visited[i]:
            visited[i] = True
            DFS(i)
    
for i in range(1, n+1):
    visited = [False] * (n + 1)
    if check[i] == 1:
        DFS(i)

print(cnt)

 

나의 풀이 (개선)

야외를 기준으로 For문을 돌리고 야외 - 실내 (야외-야외는 갯수 체크하지 않고 하나로 생각한다.)

그리고 실내 기준으로 For문을 돌려 실내만 찾았다. 

import sys
from collections import defaultdict
sys.setrecursionlimit(1000000000)
n = int(sys.stdin.readline().rstrip())
a = list(map(int, list(sys.stdin.readline().strip())))
check = [1]
for i in a:
    check.append(i)

graph = defaultdict(list)

for i in range(n-1):
    s, e = map(int, sys.stdin.readline().split())
    graph[s].append(e)
    graph[e].append(s)

def permutation(n: int, c: int) -> int:
    if c == 0:
        return 1
    return n * permutation(n-1,c-1)

def DFS_find_outin(start: int, CheckPoint: int) -> int: #실외 기준으로 탐색
    global visited
    cnt=0
    if CheckPoint == 0: # 야외와 연결된 모든 실내
        visited[start] = True
        for i in graph[start]:
            if not visited[i]:
                if check[i] == 1:
                    #visited[i] = True
                    cnt += 1
                elif check[i] == 0:
                    #visited[i] = True
                    cnt += DFS_find_outin(i, 0)
        return cnt 
    else: # CheckPoint == 1 일때, 실내 - 실내 
        for i in graph[start]:
                if check[i] == 1:
                    cnt += 1
        return cnt

visited = [False] * (n+1)
ans = 0
for i in range(1, n+1):
    if check[i] == 0 and not visited[i]: #실외인 것만 실행
        tmp = permutation(DFS_find_outin(i, 0),2)
        ans += tmp
for i in range(1, n+1):
    if check[i] == 1:
        tmp = DFS_find_outin(i, 1)
        ans += tmp
print(ans)