[백준 / BOJ] 15591번 : MooTube (Silver) - JAVA

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

15591번: MooTube (Silver)

BFS (2056ms)

각 간선을 List로 저장하고 Query가 주어질 때마다 간선 리스트를 바탕으로 BFS로 갈 수 있는 노드를 찾도록 했다. 단순한 방법으로 통과하기는 했지만 시간이 2056ms나 나왔다.

 

N개의 노드에 N - 1개의 간선이 있고, 모든 노드가 이어져 있다는 조건으로, 사이클이 없는 트리 구조라는 것을 활용해보고 싶었지만 떠오르는 방법이 없었다. 이에 순위권에 있는 다른 사람들의 풀이를 참고하여 다시 풀었다.

 

Union - Find (204ms)

다른 사람들의 풀이 중 속도가 빠른 것은 모두 union-find 방법을 사용하고 있었다. 게다가 query가 들어올 때마다 푸는 게 아니라, query를 일단 다 받고, K값이 큰 것부터 확인하면서 이전 연산을 활용했다.

 

0. 영상 간 유사도(간선)을 모두 받고 r을 기준으로 정렬한다.

나중에 query에 나오는 k값과 비교할 때, 전체를 비교할 필요 없이 필요한 부분만 비교할 수 있게 미리 연산해두는 것이다.

 

1. Query를 모두 받고 k를 기준으로 정렬한다.

어떤 Query에서의 조건은 r 값이 k값 이상일 때 카운트하는 것이다. k값이 큰 순서대로 연산을 하면 k값이 작아졌을 때 기존에 수행했던 연산의 결과를 사용할 수 있다. 단, 이 때문에 query의 순서가 바뀔 수 있으므로 각 query에 원래 index 값 또한 저장해두어야 한다.

 

2. 정렬한 Query 순서대로 다음 연산을 진행한다.
    2 - 1) 간선 리스트 안에서 간선의 r값이 k와 같거나 큰 경우 두 노드를 union한다.
    2 - 2) 조건에 맞는 간선이 더이상 없는 경우 주어진 v의 parent를 찾아 rank 값을 결과로 저장한다.

2 - 1은 v와 상관없이, r이 k와 같거나 큰 경우를 모두 이은 것이다.

위와 같은 경우에, 2 - 3, 2 - 5 사이의 r 값은 k보다 크기 때문에 union이 되지만, 1 - 2 사이의 r이 1이기 때문에 v와는 연결되지 않는다. 따라서 전부 union을 하더라도 v = 1일 때 관련 영상은 4 하나만 뜨게 되는 것이다.

 

이 다음 query가  k = 1, v = 1이라면 어떻게 될까? 그러면 1 - 2 사이 간선이 이어지면서 2, 3, 5가 한 번에 1과 연결된다. 이 연결된 노드의 개수는 rank로 관리한다. 아무 영상과도 연결이 안됐을 때는 rank가 1이고, 이후에는 연결될 때마다 부모 노드에게 두 개의 rank를 합산한 값을 저장한다.

// parent
[1, 2, 2, 1, 2]

// rank
[2, 3, 1, 1, 1]

위 그림의 상황을 나타내면 위와 같다. (보기 쉽게 idx가 낮은 것을 parent로 하도록 그렸다. 실제 코드에서는 rank 값이 더 큰 노드가 부모 노드가 되도록 했다. rank 값이 큰 것을 부모로 했을 때 find의 품이 덜 들기 때문이다.) 아까의 예시처럼 다음 query가 k = 1, v = 1이 된다면 다음과 같이 변할 것이다.

// parent
[1, 1, 1, 1, 1]

// rank
[5, 3, 1, 1, 1]

 

전체 코드

 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

// 시간 : 204ms
// 메모리 : 17888KB
public class Main {

	static int[] parent;
	static int[] rank;
	static class Edge implements Comparable<Edge>{
		int p;
		int q;
		int r;
		
		public Edge(int p, int q, int r) {
			this.p = p;
			this.q = q;
			this.r = r;
		}
		
		@Override
		public int compareTo(Edge o) {
			return o.r - this.r;
		}
	}
	
	static class Query implements Comparable<Query> {
		int idx;
		int k;
		int v;
		
		public Query(int idx, int k, int v) {
			this.idx = idx;
			this.k = k;
			this.v = v;
		}
		
		@Override
		public int compareTo(Query o) {
			return o.k - this.k;
		}
	}
	public static void main(String[] args) throws IOException {

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		
		int N = Integer.parseInt(st.nextToken());
		int Q = Integer.parseInt(st.nextToken());
		
		parent = new int[N];
		rank = new int[N];
		
		Edge[] edge = new Edge[N - 1];
		
		for (int i = 0; i < N; i++) {
			parent[i] = i;
			rank[i] = 1;
		}
		
		for (int i = 0; i < N - 1; i++) {
			st = new StringTokenizer(br.readLine());
			int p = Integer.parseInt(st.nextToken()) - 1;
			int q = Integer.parseInt(st.nextToken()) - 1;
			int r = Integer.parseInt(st.nextToken());
			
			edge[i] = new Edge(p, q, r);
		}
		
		Query[] query = new Query[Q];
		for (int i = 0; i < Q; i++) {
			st = new StringTokenizer(br.readLine());
			int k = Integer.parseInt(st.nextToken());
			int v = Integer.parseInt(st.nextToken()) - 1;
			
			query[i] = new Query(i, k, v);
		}
		
		Arrays.sort(edge);
		Arrays.sort(query);
		
		int[] res = new int[Q];
        int idx = 0;
		for (Query q : query) {
			while(idx < N - 1 && edge[idx].r >= q.k) {
				union(edge[idx].p , edge[idx].q);
				idx++;
			}
			res[q.idx] = rank[find(q.v)] - 1;
		}
		
		StringBuilder sb = new StringBuilder();
		for (int i = 0; i < Q; i++) {
			sb.append(res[i]).append("\n");
		}
		System.out.println(sb);
	}
	
	private static void union(int p, int q) {
		
		int pp = find(p);
		int pq = find(q);
		
		if(pp == pq) return;
        if(rank[pp] >= rank[pq]) {
            parent[pq] = pp;
    		rank[pp] += rank[pq];
        } else {
            parent[pp] = pq;
            rank[pq] += rank[pp];
        }
		
	}

	private static int find(int p) {
		if(parent[p] == p) return p;
		return parent[p] = find(parent[p]);
	}

}