백준 25402 트리와 쿼리 파이썬 풀이

이 문제는 서브테스크를 긁고 최적화를 할 방법을 모르겠어서 풀이를 봤다. 풀이가 신기하고 좋은 문제라고 생각해서 풀이를 올려본다.

 

https://www.acmicpc.net/problem/25402

25402번: 트리와 쿼리

이 문제를 요악하자면 트리를 S라는 집합의 정점들로만 한정지었을때, 생기는 각 컴포넌트들의 연결 되어있는 노드의 개수를 구하면 된다. 연결 개수는 컴포넌트의 노드 개수를 n이라 했을때, nC2가 된다. 

나이브하게 매 쿼리마다 dfs를 돌리며 컴포넌트의 개수를 세어주면, $O(NQ)$가 된다. 이로 서브테스크 3까지 긁을 수 있다 (24점 풀이).

import sys
input = sys.stdin.readline

def find(x):
    if x == p[x]: return x
    p[x] = find(p[x])
    return p[x]

def union(x, y):
    x, y = find(x), find(y)
    if x == y: return
    if r[x] < r[y]: x, y = y, x
    p[y] = x
    r[x] += r[y]

def dfs(n):
    if visited[n]: return 0
    visited[n] = True
    for c in g[n]:
        if not visited[c] and c in s:
            union(n, c)
            dfs(c)
    return r[n]

n = int(input())
g = [[] for _ in range(n + 1)]
for _ in range(n-1):
    u, v = map(int, input().split())
    g[u].append(v)
    g[v].append(u)

Q = int(input())
for _ in range(Q):
    p = [i for i in range(n + 1)]
    r = [1 for i in range(n + 1)]
    k, *s = map(int, input().split())
    visited = [0] * (n + 1)
    s = set(s)
    ans = 0
    for node in s:
        tmp = dfs(node)
        ans += (tmp * (tmp - 1)) // 2
    print(ans)

여기서 최적화를 해주면 만점 풀이를 얻을 수 있다. 핵심적인 아이디어는 트리의 성질이다. 임이의 정점을 루트로 잡았을때, 각 노드의 부모는 유일하게 존재한다. 나이브하게 dfs를 돌리면 S의 각 정점마다 그 정점의 차수만큼 탐색을 해주어야하는데, 각 노드들의 부모를 저장해주면, 노드의 부모가 S에 있는지만 확인해주면 된다. 유니온파인드를 이용해서 컴포넌트의 개수를 세는것은 쉽다. 즉, 각 쿼리를 $O(K * \alpha(K))$에 해결해줄 수 있고, K의 합은 1e6 이하이기에 이는 충분한 시간복잡도이다.

import sys
input = sys.stdin.readline
print = sys.stdout.write

def find(x):
    while x != p[x]:
        p[x] = p[p[x]]
        x = p[x]
    return p[x]

def union(x, y):
    x, y = find(x), find(y)
    if x == y: return
    if r[x] < r[y]: x, y = y, x
    p[y] = x
    r[x] += r[y]

def dfs(n):
    visited[n] = 1
    s = [n]
    while s:
        v = s.pop()
        for w in g[v]:
            if not visited[w]:
                an[w] = v
                s.append(w)
                visited[w] = 1

n = int(input())
n += 1
g = [[] for _ in range(n)]
for _ in range(n-2):
    u, v = map(int, input().split())
    g[u].append(v)
    g[v].append(u)

an = [0] * n
visited = [0] * n
dfs(1)


p = [i for i in range(n)]
r = [1 for i in range(n)]
isin = [0] * n
Q = int(input())
for _ in range(Q):
    ans = 0
    k, *s = map(int, input().split())
    for i in range(k): isin[s[i]] = 1
    for node in s:
        if isin[an[node]]: union(node, an[node])
    for node in s:
        if p[node] == node: ans += (r[node] * (r[node] - 1)) // 2
        p[node] = node
        r[node] = 1
    for i in range(k): isin[s[i]] = 0
    print(str(ans) + '\n')

dfs로 각 노드의 부모를 전처리해주고, 유니온파인드를 통해 컴포넌트에 속한 노드의 개수를 세주면 된다. 파이썬 사용자는 주의해야할 점이 있는데, dfs와 유니온파인드를 재귀 구현하면 시간초과가 난다. 전부 비재귀로 바꾸고 나서야 결국 AC를 받을 수 있었다.