파이썬 백준 21606 아침 산책

import sys input = sys.stdin.readline sys.setrecursionlimit(10**8) n = int(input()) C = list(input()) del C[-1] out = [] for i in range(n) :     C[i] = int(C[i])     if C[i] == 0 :         out.append(i+1) count = 0 V = [[] for i in range(n + 1)] for i in range(n - 1) :     v1, v2 = map(int, input().split())     if C[v1 - 1] + C[v2 - 1] == 2 :         count = count + 2     else :         V[v1].append(v2)         V[v2].append(v1) sum = sum(C) if sum == 0 :     print(0) elif sum == 1 :     print(0) elif sum == 2 :     print(2) else :     def DFS(before, now) :         global ct         if C[now - 1] == 1 :             ct = ct + 1         else :             C[now - 1] = -1             for next in V[now] :                 if next != before :                     DFS(now, next)     for i in out :         if C[i-1] == 0 :             ct = 0             DFS(0, i)             count = count + ct * (ct - 1)     print(count)

처음 접근은 저 실외 0, 실내 1값을 코스트로 해서 코스트가 2가 되면 종료하며 카운트하는 DFS였다.

근데 108점이 뜨길래 몇번의 최적화를 거쳤다.

sum(C)가 0, 1, 2면 계산 없이 바로 자명해를 뱉는 저 if문들이 그 흔적이다.

근데 아무리 줄여도 시간초과로 200점이 안나와서 고민하고 있었는데, 이번에도 정답자 형님의 힌트를 얻었다.

 

서로 연결된 실외를 하나의 영역으로 보았을 때, 그 영역에 연결된 실내의 갯수가 n개라면 그 덩어리(?)로부터 만들 수 있는 산책 경로의 수는 nP2이다.

또한 서로 연결된 두 실내로부터는 실외를 거치지 않는 2개의 산책 경로가 만들어진다.

즉 이 둘을 합치면 총 개수가 나온다.

 

v1, v2 연결 정보를 받을 때 두 점이 모두 실내인지 확인하여 count를 +2하는 것으로 후자의 경우를 먼저 센다.

그 후 연결된 점으로 확산하며 0이면 -1로 만들고(= 방문기록을 남기고) 1이면 ct를 세는 DFS를 정의한다.

실외 점을 모아둔 리스트에서, 아직 0인 점(이전 점의 DFS에서 방문기록이 안남은 점)에 대해 DFS를 실행해주며 count를 ctP2만큼 더한다.

 

이 과정을 모두 거친 후의 count가 총 산책 경로의 개수가 된다.