BOJ 백준 16453 Linhas de Metrô

 

문제 : https://www.acmicpc.net/problem/16453

[ 문제 요약 ]

• 노드의 개수와 질의 개수가 주어집니다
• 노드 개수 - 1개의 간선이 있는 트리 정보가 주어지며, 트리 정보는 두 노드 번호로 주어집니다.
• 질의는 a, b, c, d로 되어있으며, a-b 구간과 c-d 구간의 겹치는 노드의 개수를 출력하는 문제입니다.

 

[ 테스트 케이스 설명 ]

10 4	// 노드 개수(5<=100,000), 질의 수(1<=20,000)
1 4	// 이어진 두 노드 번호
4 5
3 4
3 2
7 3
6 7
7 8
10 8
8 9
6 10 2 5	// (A,B), (C,D) 노드가 주어짐
1 9 5 10
9 10 2 1
5 10 2 9
//각 질의마다 두 노드가 서로 공유하는 간선을 출력
0
4
0
3

 

[ 문제 해설 ]

 

이 문제는 노드 수가 10만이고, 질의가 2만이기 때문에 완전 탐색으로는 풀 수 없습니다.

 

노드에서 노드까지 경로를 한 번에 찾아야 하고, 추가로 해당 노드가 겹치는지까지 알아야 하기 때문에 느리게 갱신되는 세그먼트 트리와 HLD 알고리즘을 사용해야 합니다.

 

트리상에서 HLD는 두 노드 사이의 경로를 정확히 한 번에 찾아갈 수 있도록 하고,

 

느리게 갱신되는 세그먼트 트리는, 구간 업데이트를 할 때 사용됩니다. 구간 업데이트가 필요한 이유는 아래와 같습니다.

 

A-B 노드로 갈 때 해당 구간에 있는 노드들에 모두 1을 더합니다.

 

그리고 C-B 노드를 탐색할 때 그 사이 경로에 있는 노드들 값의 합을 구간 합 쿼리로 구해옵니다.

 

그러면 겹치는 노드가 뭔지 알 수 있게 됩니다.

 

마지막으로 1로 업데이트했던 A-B 구간의 노드 값들을 다시 -1을 더 해주어 모든 노드 값을 0으로 복구시켜줍니다.

 

HLD 알고리즘에 대략적인 설명은 아래 문제에서 참고 부탁드립니다.

https://kyjdummy.tistory.com/21

 

디테일은 코드를 통해 확인 가능합니다.

 

[ 정답 코드 ]

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.StringTokenizer;
class Main{
	
	static int N, Q;
	static int idx;
	static int [] treeIdx;
	static int [] tree;
	static int [] lazy;
	static int [] chainHead;
	static int [] chainLevel;
	static int [] chainParent;
	static ArrayList<Integer>[] adNode;
	
	public static void main(String[] args)throws Exception{
		BufferedReader	br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		StringBuilder	sb = new StringBuilder();
		N		= Integer.parseInt(st.nextToken());// 노드 개수(5<=100,000)
		Q		= Integer.parseInt(st.nextToken());// 질의 수(1<=20,000)
		treeIdx		= new int[N + 1];
		tree		= new int[N * 4];
		lazy		= new int[N * 4];
		chainHead	= new int[N + 1];
		chainLevel	= new int[N + 1];
		chainParent = new int[N + 1];
		adNode		= new ArrayList[N + 1];
		
		for(int i=0; i<=N; i++)
			adNode[i] = new ArrayList<>();
		
		for(int i=1; i<N; i++)
		{
			st = new StringTokenizer(br.readLine());
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			adNode[a].add(b);
			adNode[b].add(a);
		}
		// 해당 함수에서 자식 노드가 큰것을 인접리스트에 가장 앞으로 옮김
		setSize(1, 0, new int[N + 1]);
		
		chainHead[1] = 1;// 1번 노드의 헤드는 자기자신
		// chain정보 입력
		setHLD(1, 0, 1);
		
		while(Q-->0)
		{
			st = new StringTokenizer(br.readLine());
			// (A,B), (C,D) 노드가 주어짐
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			int c = Integer.parseInt(st.nextToken());
			int d = Integer.parseInt(st.nextToken());
			
			rangeUpdate(a, b, 1);	// a-b구간을 탐색하며 해당 구간에 1을 업데이트
			
			int ans = getAns(c, d);	// c-d구간을 탐색하며 구간 sum을 구해서 1인 노드들의 합을 구해옴
			
			rangeUpdate(a, b, -1);	// a-b구간을 탐색하며 해당 구간에 -1로 업데이트
			
			sb.append(ans).append('\n');
		}
		System.out.print(sb);
	}
	static int getAns(int n1, int n2) {
		int ans = 0;
		// chainHead가 같아질 때 까지 레벨을 낮춤
		while(chainHead[n1] != chainHead[n2])
		{
			if(chainLevel[n1] > chainLevel[n2])
			{
				int tmp = n1;
				n1 = n2;
				n2 = tmp;
			}
			ans += query(1, 1, N, treeIdx[chainHead[n2]], treeIdx[n2]);
			n2 = chainParent[n2];
		}
		// 같은 헤드일 때 마지막 연산
		if(treeIdx[n1] > treeIdx[n2])
		{
			int tmp = n1;
			n1 = n2;
			n2 = tmp;
		}
		ans += query(1, 1, N, treeIdx[n1], treeIdx[n2]);
		
		return ans;
	}
	static void rangeUpdate(int n1, int n2, int value) {
		// chainHead가 같아질 때 까지 레벨을 낮춤
		while(chainHead[n1] != chainHead[n2])
		{
			if(chainLevel[n1] > chainLevel[n2])
			{
				int tmp = n1;
				n1 = n2;
				n2 = tmp;
			}
			update(1, 1, N, treeIdx[chainHead[n2]], treeIdx[n2], value);
			n2 = chainParent[n2];
		}
		// 같은 헤드일 때 마지막 연산
		if(treeIdx[n1] > treeIdx[n2]) {
			int tmp = n1;
			n1 = n2;
			n2 = tmp;
		}
		update(1, 1, N, treeIdx[n1], treeIdx[n2], value);
	}
	static void update(int treeNode, int s, int e, int left, int right, int value) {
		propagate(treeNode, s, e);
		if(e < left || right < s)
			return;
		if(left<=s && e<= right)
		{
			lazy[treeNode] = value;
			propagate(treeNode, s, e);
			return;
		}
		int mid = (s + e) >> 1;
		
		update(treeNode << 1, s, mid, left, right, value);
		update(treeNode << 1 | 1, mid + 1, e, left, right, value);
		
		tree[treeNode] = tree[treeNode << 1] + tree[treeNode << 1 | 1];
	}
	static int query(int treeNode, int s, int e, int left, int right) {
		propagate(treeNode, s, e);
		if(e < left || right < s)
			return 0;
		if(left<=s && e<=right)
			return tree[treeNode];
		
		int mid = (s + e) >> 1;
		
		int l = query(treeNode << 1, s, mid, left, right);
		int r = query(treeNode << 1 | 1, mid + 1, e, left, right);
		
		return l + r;
	}
	static void propagate(int treeNode, int s, int e) {
		if(lazy[treeNode] != 0)
		{
			tree[treeNode] += (e - s + 1) * lazy[treeNode];
			if(s != e)
			{
				lazy[treeNode << 1]		+= lazy[treeNode];
				lazy[treeNode << 1 | 1] += lazy[treeNode];
			}
			lazy[treeNode] = 0;
		}
	}
	static void setSize(int nowNode, int parentNode, int[] size) {
		int heavySize = 0;
		int heavyIdx = 0;
		size[nowNode] = 1;
		for(int i=0; i<adNode[nowNode].size(); i++)
		{
			int next = adNode[nowNode].get(i);
			
			if(next == parentNode)
				continue;
			
			setSize(next, nowNode, size);
			
			size[nowNode] += size[next];
			if(heavySize < size[next])
			{
				heavySize = size[next];
				heavyIdx = i;
			}
		}
		if(adNode[nowNode].size() > 0)
			Collections.swap(adNode[nowNode], 0, heavyIdx);
	}
	static void setHLD(int nowNode, int parentNode, int level) {
		treeIdx[nowNode] = ++idx;
		chainLevel[nowNode] = level;
		
		for(int i=0; i<adNode[nowNode].size(); i++)
		{
			int next = adNode[nowNode].get(i);
			
			if(next == parentNode)
				continue;
			
			if(i == 0)	// heavy 노드일 경우 체인 유지
			{
				chainHead[next] = chainHead[nowNode];
				chainParent[next] = chainParent[nowNode];
				setHLD(next, nowNode, level);
			}
			else		// light일 경우 새 체인 생성
			{
				chainHead[next] = next;		// 체인의 헤드는 자기자신
				chainParent[next] = nowNode;// 이전 체인으로 바로 점프를 위한 이전 노드 값 삽입
				setHLD(next, nowNode, level + 1);// 새 체인시 레벨을 올린다.
			}
		}
	}
}