BOJ 백준 13309 트리

 

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

[ 문제 요약 ]

• 트리가 주어지고, 두 정점이 주어집니다. 이때 단순히 두 정점이 이어져 있는지 체크하는 쿼리와, 이어져 있거나 이어져 있지 않을 경우를 파악하고, 이어져 있다면 먼저 입력된 노드를 그 부모노드와 끊고, 이어져있지 않다면 나중에 입력된 노드를 그 노드의 부모 노드와 끊습니다.

 

[ 테스트 케이스 설명 ]

3 4	// 트리의 정점개수(1 ≤ 200,000), 질의 개수(1 ≤ 200,000)
1		// 부모 정점 번호(1<=N)
1		// 부모 정점 번호(1<=N)
2 3 1	// b,c,d / d가1이면 b와 c를 연결하는 경로가 있는지 파악후 있으면 b의 부모 정점과 b연결 엣지 제거, 없으면, c의 부모정점과 c를 연결하는 에지 제거
1 3 0	// b,c,d / d가0이면 b와 c를 연결하는 경로가 있는지만 파악
2 3 1
1 3 1
//답
YES
YES
NO
NO

[ 문제 해설 ]

• 이 문제는 HLD 알고리즘을 쉽게 풀 수 있습니다. HLD 알고리즘은 아래 링크의 [문제 해설]을 보시면 설명되어 있습니다.
• https://kyjdummy.tistory.com/21
• 문제에서 요구하는 것은 두 정점이 연결되어 있는지, 아닌지 확인하는 것이고, 또한 연결을 끊는 기능도 필요로 합니다.
• 이를 구현하기 위해 boolean 세그먼트 트리를 만듭니다. 연결되어 있다면 false, 연결이 끊기면 true로 하여, 두 노드 사이에서 세그먼트 트리 쿼리 연산 결과가 true가 나오면 두 노드는 끊어져있는 것으로 보면 됩니다.
• d 값이 1인 경우 연결 유무에 따라 끊어야 하는 노드가 다르기 때문에 분기를 잘해야 합니다.
• 그리고 두 노드의 간선 정보를 세그먼트 트리에 저장해야 합니다. 부모 노드, 자식 노드가 있을 때, 간선의 정보는 자식 노드에 저장해야 합니다.
• 디테일은 코드를 통해 확인 가능합니다.

 

 

[ 정답 코드 ]

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[] chainHead;
	static int[] chainLevel;
	static int[] chainParent;
	static boolean[] tree;
	static ArrayList<Integer>[] adNode;

	public static void main(String[] ags)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());	// 트리의 정점개수(1 ≤ 200,000)
		Q			= Integer.parseInt(st.nextToken());	// 질의 개수(1 ≤ 200,000)
		tree		= new boolean[N * 4];				// 초기값 다 false, 연결시 false, 연결을 끊으면 true
		treeIdx		= new int[N + 1];
		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 child=2; child<=N; child++)
		{
			int parent = Integer.parseInt(br.readLine());
			adNode[parent].add(child);// 단방향 연결
		}
		
		// 해당 함수에서 HLD를 위한 무거운 자식 노드를 맨앞으로 옮긴다.
		moveHeavyChildToFront(1, 0, new int[N + 1]);
		
		chainHead[1] = 1;	// 1번노드의 헤드는 자기 자신
		
		setHLD(1, 0, 1);	// HLD 세팅
		
		while(Q-->0)
		{
			st = new StringTokenizer(br.readLine());
			int b = Integer.parseInt(st.nextToken());
			int c = Integer.parseInt(st.nextToken());
			int d = Integer.parseInt(st.nextToken());
			boolean isPossible = true;
			
			int node1 = b;
			int node2 = c;
			// chainHead가 같아질 때 까지, 즉, 같은 체인에 속할 때 까지 반복
			while(isPossible && chainHead[node1] != chainHead[node2])
			{
				// 레벨이 낮은 것(값이 큰것)을 node2로 바꿔서 낮은 레벨을 지속적으로 위로 올린다. 
				if(chainLevel[node1] > chainLevel[node2])
				{
					int tmp = node1;
					node1 = node2;
					node2 = tmp;
				}
				
				boolean notConnected = query(1, 1, N, treeIdx[chainHead[node2]], treeIdx[node2]);
				
				if( !notConnected )				// 노드가 연결되어있는게 맞다면
					node2 = chainParent[node2];	// 다음 체인으로 점프
				else
					isPossible = false;
			}
			
			// 위 코드들을 통해 같은 체인에 속하게 되었을 것이고, 연결 유무를 최종 확인
			if(isPossible)
			{
				if(treeIdx[node1] > treeIdx[node2])
				{
					int tmp = node1;
					node1 = node2;
					node2 = tmp;
				}
				
				boolean notConnected = query(1, 1, N, treeIdx[node1] + 1, treeIdx[node2]);
				if(notConnected)
					isPossible = false;
			}
			
			if(d == 1)// d가 1이면
			{
				int deleteNode = c;	// b와 c를 연결하는 통로가 없으면 c의 부모정점과 c를 연결하는 엣지 제거
				
				if(isPossible)		// b와 c를 연결하는 통로가 있으면 b의 부모 정점과 b연결 엣지 제거
				{
					deleteNode = b;
				}
				
				update(1, 1, N, treeIdx[deleteNode], true);
			}
			
			sb.append(isPossible ? "YES\n" : "NO\n");
		}
		System.out.print(sb);
	}

	static boolean query(int treeNode, int s, int e, int left, int right) {
		
		if(e < left || right < s)
			return false;

		if(left<=s && e<=right)
			return tree[treeNode];
		
		int mid = (s + e) >> 1;
		
		return query(treeNode << 1, s, mid, left, right)
				| query(treeNode << 1 | 1,mid + 1, e, left, right);
	}
	static void update(int treeNode, int s, int e, int targetIdx, boolean value) {
		
		if(e < targetIdx || targetIdx < s)
			return;
		
		if(s == e)
		{
			tree[treeNode] = value;
			return;
		}
		
		int mid = (s + e) >> 1;
		
		update(treeNode << 1, s, mid, targetIdx, value);
		update(treeNode << 1 | 1, mid + 1, e, targetIdx, value);
		
		tree[treeNode] = tree[treeNode << 1] | tree[treeNode << 1 | 1];
	}
	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 child = adNode[nowNode].get(i);
			
			if(child == parentNode)
				continue;
			
			if(i == 0)	// heavy인 경우, 체인 그대로 유지
			{
				chainHead[child]	= chainHead[nowNode];	// heavy는 부모의 값을 그대로 저장
				chainParent[child]	= chainParent[nowNode];	// heavy는 부모의 값을 그대로 저장
				setHLD(child, nowNode, level);
			}
			else		// light인 경우, 새로운 체인 시작
			{
				chainHead[child]	= child;				// head는 자기자신
				chainParent[child]	= nowNode;				// 다음 점프할 노드는 자기의 부모노드로 세팅
				setHLD(child, nowNode, level + 1);			// 레벨을 늘려 새로운 체인 시작
			}
		}
	}
	static void moveHeavyChildToFront(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 child = adNode[nowNode].get(i);
			
			if(child == parentNode)
				continue;
			
			moveHeavyChildToFront(child, nowNode, size);
			
			size[nowNode] += size[child];
			
			if(heavySize < size[child])
			{
				heavySize = size[child];
				heavyIdx = i;
			}
		}
		
		if(adNode[nowNode].size() > 0)
			Collections.swap(adNode[nowNode], 0, heavyIdx);
	}
}