[Baekjoon] 백준 15591 - MooTube (Silver) by 파이썬 Python

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

 

문제 설명

그래프와 유니온-파인드를 활용해 특정 조건을 만족하는 노드(동영상)들의 수를 구하는 문제이다.

 

그래프 구성

• 동영상들은 노드, 동영상 쌍 간의 USADO 값은 간선의 가중치로 주어진다.
• 총 N개의 노드와 N-1개의 간선이 있으며, 모든 노드는 하나의 연결 그래프를 이룬다.

 

USADO 경로 계산: 두 노드 간의 USADO는 그들을 연결하는 경로의 간선 중 최솟값

 

아이디어

처음 풀이로 BFS를 사용해 탐색했다. 하지만 시간초과가 발생했다.

 

따라서 union-find 를 사용해 해결했다. 또한 간선을 내림차순 정렬하여 K 이상의 간선만 처리하도록하여 효율성을 높였다.

 

코드

노드 u가 속한 집합의 루트노드를 찾는 find 함수이다. 

재귀함수를 활용해 부모 노드를 찾고 저장하여 반환한다.

def find(u):
    if parent[u] != u:
        parent[u] = find(parent[u])
    return parent[u]

 

 

노드 U와 V가 속한 두 집합을 합치는 union 함수는 아래와 같다.

집합의 크기를 활용하여 작은 집합을 큰 집합에 합친다.

def union(u, v):
    u_root = find(u)
    v_root = find(v)

    if u_root != v_root:
        if size[u_root] < size[v_root]:
            parent[u_root] = v_root
            size[v_root] += size[u_root]
        else:
            parent[v_root] = u_root
            size[u_root] += size[v_root]

 

 

간선에 대한 입력 질의를 처리할 때는 K값 이상인 간선을 union-find에 추가하여 계산하도록 하였다.

result = [0] * q
edge_idx = 0

for k_value, v_node, idx in queries:
    while edge_idx < len(edges) and edges[edge_idx][0] > k_value:
        usado_value, u_node, v_node = edges[edge_idx]
        union(u_node, v_node)
        edge_idx += 1
    
    root_v = find(v_node)
    result[idx] = size[root_v] - 1

 

 

전체 코드는 아래와 같다.

import sys
sys.setrecursionlimit(1 << 25)
def find(u):
    if parent[u] != u:
        parent[u] = find(parent[u])
    return parent[u]

def union(u, v):
    u_root = find(u)
    v_root = find(v)
    if u_root != v_root:
        if size[u_root] < size[v_root]:
            parent[u_root] = v_root
            size[v_root] += size[u_root]
        else:
            parent[v_root] = u_root
            size[u_root] += size[v_root]

n, q = map(int, sys.stdin.readline().split())
edges = []
for _ in range(n - 1):
    p, node, r = map(int, sys.stdin.readline().split())
    edges.append((r, p, node)) 

queries = []
for i in range(q):
    k, v = map(int, sys.stdin.readline().split())
    queries.append((k, v, i)) 

edges.sort(reverse=True)
queries.sort(reverse=True)

parent = [i for i in range(n + 1)]
size = [1] * (n + 1)

result = [0] * q
edge_idx = 0
for k_value, v_node, idx in queries:
    # 현재 K 값 이상인 간선을 유니온 파인드에 추가
    while edge_idx < len(edges) and edges[edge_idx][0] >= k_value:
        usado_value, u_node, v_node_edge = edges[edge_idx]
        union(u_node, v_node_edge)
        edge_idx += 1
    root_v = find(v_node)
    result[idx] = size[root_v] - 1  # 시작 노드를 제외한 연결된 노드의 수

for res in result:
    print(res)