[알고리즘] BOJ 15681 트리와 쿼리 #Python

[BOJ] 15681 트리와 쿼리 바로가기

📍 문제

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

• 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

---

📍 입력

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)

이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)

이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.

이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)

입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

---

📍 출력

Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

---

📍 풀이

💡 고찰

• 보자마자 Union & Find 알고리즘을 이용한 방법으로 문제에 접근하게 되었다.
• 문제를 더 읽어보니 하단의 힌트 에서 주어진 makeTree() 함수와 countSubtreeNodes() 함수를 코드에 적용하여 문제를 해결하는 방식이었다.
• 다른 사람이 구현해 놓은 코드를 그대로 사용하는 것이 간단한 문제인 줄 알았는데, 문제의 의도에 맞게 적용하여 사용하는 것은 생각보다 어려웠다.
• 마치 처음 API를 사용할 때 느꼈던 경험이 되살아나는 느낌이었다.

📌 문제 풀이

✏️ 각 정점 연결 관계 설정

• defaultdict(list) 로 정의한 connect 딕셔너리를 생성하여 각 정점의 연결 관계 정보를 저장한다.

connect = defaultdict(list)
for u, v in UV:
    connect[u].append(v)
    connect[v].append(u)

✏️ makeTree() 함수

• currentNode 를 루트로 하는 서브트리를 생성한다.
  • currentNode 에 연결된 모든 자식 정점들을 순회한다.
    • 만약 자식 정점(Node)과 부모 정점(parent)이 같지 않다면,
      • tree의 currentNode 리스트에 현재 정점(Node)를 추가한다.
      • 현재 정점(Node)의 부모를 currentNode 로 설정한다.
      • 현재 정점(Node)을 루트로 하는 makeTree() 함수를 실행한다.

def makeTree(currentNode, parent) :
    for Node in connect[currentNode] :
        if Node != parent:
            #add Node to currentNode’s child
            tree[currentNode].append(Node)
            #set Node’s parent to currentNode
            parents[Node] = currentNode
            makeTree(Node, currentNode)

✏️ countSubtreeNodes() 함수

• 현재 정점(currentNode)을 루트로 하는 서브트리에 속한 정점의 수를 센다.
  • 초기값은 자기 자신도 포함해야 하므로 1로 초기화한다.
  • 현재 정점의 자식 정점들을 순회한다.
    • 자식 정점(Node)를 인자로 countSubtreeNodes() 함수를 실행한다.
    • 모든 자식 정점에 대한 계산이 완료되었으면, 현재 정점에 자식노드의 수를 모두 더한다.

def countSubtreeNodes(currentNode) :
        size[currentNode] = 1 # 자신도 자신을 루트로 하는 서브트리에 포함되므로 0이 아닌 1에서 시작한다.
        for Node in tree[currentNode]:
            countSubtreeNodes(Node)
            size[currentNode] += size[Node]

✍ 코드

from sys import stdin, setrecursionlimit
from collections import defaultdict

setrecursionlimit(10**6) # 최대 재귀 깊이의 값을 10^6 으로 재설정

def solution(N, R, UV, U):

    connect = defaultdict(list)
    parents = list(range(N+1))
    tree = defaultdict(list)
    size = list(0 for _ in range(N+1))

    # 연결 관계 초기화
    for u, v in UV:
        connect[u].append(v)
        connect[v].append(u)

    # 각 정점을 루트로 하는 서브트리에 속한 정점의 수 count
    def countSubtreeNodes(currentNode) :
        size[currentNode] = 1 # 자신도 자신을 루트로 하는 서브트리에 포함되므로 0이 아닌 1에서 시작한다.
        for Node in tree[currentNode]:
            countSubtreeNodes(Node)
            size[currentNode] += size[Node]

    # R을 정점으로 하는 트리 생성
    def makeTree(currentNode, parent) :
        for Node in connect[currentNode] :
            if Node != parent:
                #add Node to currentNode’s child
                tree[currentNode].append(Node)
                #set Node’s parent to currentNode
                parents[Node] = currentNode
                makeTree(Node, currentNode)
        

    makeTree(R,-1)
    countSubtreeNodes(R) 

    # 정점 u를 루트로 하는 서브트리에 속한 정점의 수 출력
    for u in U:
        print(size[u])

N, R, Q = map(int,stdin.readline().split())
UV = list(list(map(int,stdin.readline().split())) for _ in range(N-1))
U = list(int(stdin.readline()) for _ in range(Q))

solution(N, R, UV, U)