[Python] 15586. MooTube (Gold)

문제

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

15586번: MooTube (Gold)

 

---

풀이 과정

동영상과 유사도를 가지고 그래프를 만들어 농부 존이 제시한 동영상 v에서 유사도 k 이상으로 이어진 동영상이 몇 개 있는지 묻는 문제이다.
여기서 주의할 점은 존은 두 쌍 사이의 동영상의 유사도를 그 경로의 모든 연결들의 유사도 중 최솟값으로 한다.

MooTube (Silver) (BOJ 15591)에서는 각 노드마다 BFS를 돌려 통과했지만, 이번에는 더 큰 높은 난이도의 입력 제한 때문에 더 효율적인 풀이를 떠올려야 했다.
스터디 팀원이 입력받은 간선정보와 쿼리를 정렬하고 union & find를 사용하여 시간 복잡도를 획기적으로 줄여 통과하였는데, 이 놀라운 풀이를 복습하고 잊지 않고자 블로그에 올리게 되었다. 

문제 풀이의 순서는 다음과 같다.

1. 입력받은 간선 정보를 유사도를 기준으로 내림차순 정렬한다.
2. 입력받은 쿼리를 K를 기준으로 내림차순 정렬한다.
   * 정렬 하는 순간 순서에 대한 개념이 사라져서 답을 순서대로 낼 수 없기 때문에 정렬 전에 쿼리마다 고유의 인덱스를 부여한다.
3. 정답을 담을 배열과 size배열을 초기화한다.
   * size 배열은 각 노드를 루프노드로 하는 트리에 연결된 자손 노드 수를 저장하는 역할을 한다.
   * size 배열의 모든 값은 1로 초기화, union을 할때마다 size 배열의 값을 업데이트
4. 쿼리와 간선정보를 순회한다.
   4-1. 현재 쿼리에서 간선 정보 배열을 순회하며 유사도가 K보다 큰 간선을 모두 union 한다.
   4-2. 순회가 끝나면 쿼리에서 물어본 노드번호가 속한 트리(union 된 집합을 의미)의 노드수를 정답 배열에 저장한다.
           (find로 해당트리의 루트 노드를 찾은 후 size배열의 루트 노드 인덱스를 찾는다.)
   4-3. 모든 쿼리를 순회한 후 정답배열에 저장된 값을 순서대로 출력

 

최적화 방법

1. 쿼리배열이 K에 대한 내림차순으로 정렬되어 있기 때문에 다음 쿼리에서는 이전에 union 한 간선을 추가로 볼 필요가 없다.
   즉, 쿼리를 순회하며 union이 되지 않은 간선만 추가로 union하면 된다.
2. 간선 정보 배열이 유사도 기준 내림차순 정렬되어 있기 때문에 현재 쿼리의 K보다 작은 유사도를 가진 간선 정보를 만날 시,
   그때의 인덱스를 임시로 저장해놓고 정답을 저장한 후에 다음 쿼리에서 저장한 인덱스부터 다시 순회를 시작한다.
   이렇게 하면 모든 쿼리와 간선을 한번씩만 보게 되므로, O(n + q)의 시간 복잡도를 가지게 된다.(노드 수 n, 쿼리 수 q)

 

---

코드

import sys
sys.setrecursionlimit(100000)
input = sys.stdin.readline

def find(x):
    if par[x] == x:
        return x
    else:
        par[x] = find(par[x])
        return par[x]

def union(x, y):
    x = find(x)
    y = find(y)

    if x == y:
        return

    par[x] = y
    size[y] += size[x]

n, q = map(int, input().split()) # 노드 수, 쿼리 수
v = [list(map(int, input().split())) for i in range(n - 1)] # 연결 정보
query = [list(map(int, input().split())) + [i] for i in range(q)] # 쿼리 정보들, 쿼리 고유 인덱스 부여
ans = [0 for i in range(q)]

v.sort(key=lambda x:x[2], reverse=True) # v는 k기준 내림차순
query.sort(reverse=True) # 쿼리는 k기준 내림차순
par = list(range(n + 1))
size = [1 for i in range(n + 1)]


idx = 0
for i in query: 
    while idx < len(v) and v[idx][2] >= i[0]:
        union(v[idx][0], v[idx][1])
        idx += 1

    ans[i[2]] = size[find(i[1])] - 1 # 쿼리 인덱스에 답저장

for i in range(q):
    print(ans[i])