백준 21606: 아침 산책 (Python) 참고용

✋ 본 글은 제 독창적인 알고리즘 풀이가 아니며, 파이썬으로 설명된 자료가 구글링을 해도 나오지 않아 kth990303 님의 아이디어 설명과 C++로 올려주신 풀이를 참고해서 파이썬으로 짠 코드를 공유하고자 합니다. 아이디어 설명은 kth990303의 코딩 블로그 링크를 참고해주시기 바라며, 제 코드에 대한 설명은 코드에 적힌 주석으로 갈음합니다.

백준 21606: 아침산책 문제 링크

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**9)

def calPaths(graph: list, col: list) -> int:
    count = 0
    visited = set()

    def dfs(exterior: int) -> int:
        cnt = 0
        for neighbor in graph[exterior]:
            if col[neighbor] == 1:
                cnt += 1
            else:
                if neighbor not in visited:
                    visited.add(neighbor)
                    cnt += dfs(neighbor)
        return cnt

    for i in range(1, numVertices + 1):
        # 각 실내별 인접한 실내 구하기
        if col[i] == 1:
            for j in graph[i]:
                if col[j] == 1:
                    count += 1
        # 인접한 실외를 한 덩어리로 보고 그 덩어리에 인접한 실내의 수를 구한 뒤 
        # 각 덩어리별로 n*(n-1)의 경우의 수를 계산
        else:
            if i not in visited:
                visited.add(i)
                tmp = dfs(i)
                count += tmp * (tmp - 1)
 
    return count

if __name__ == '__main__':
    numVertices = int(input())
    col = list(map(int, list("0"+input().strip())))

    graph = [[] for _ in range(numVertices + 1)]
    
    for _ in range(1, numVertices):
        v1, v2 = map(int, input().split())
        graph[v1].append(v2)
        graph[v2].append(v1)

    print(calPaths(graph, col))