[백준 1068번][Java] 트리 - 리프노드 세기, 노드 지우기

www.acmicpc.net/problem/1068

1068번: 트리

 

단순히 DFS 한 번에 풀 수 있는 문제인 줄 알았지만, 처리해야될 예외 사항이 좀 있습니다.

 

각 서브트리 루트에 리프노드의 개수를 저장하는 알고리즘이 있다고 합시다. 

그 알고리즘을 이용해 위와 같은 결과를 만들어 냈습니다. 여기서 2번 노드를 지우려면 루트 노드의 값(3)에서 2번 노드의 값(1)을 빼면 답이 됩니다.

 

하지만... 다음의 경우는 어떨까요?

 

위에서 1번 노드를 제거하면 값이 0이 되어버립니다. 이처럼 해당 노드를 제거하면 루트 자신이 리프노드가 되는 경우 문제가 생깁니다. 위 예외들을 신경쓰지 않는 확실한 방법은 먼저 해당 노드를 트리에서 제외시키고 리프노드를 세는 겁니다.

 

 

import java.io.*;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;

/*
트리
https://www.acmicpc.net/problem/1068
 */
public class Main {
	static int n, removeNode;
	static Node[] nodes;
	public static void main(String[] args) throws IOException {
//		InputReader reader = new InputReader("testcase.txt");
		InputReader reader = new InputReader();
		n = reader.readInt();
		nodes = new Node[n];
		for (int i = 0; i < n; i++) {
			nodes[i] = new Node(i);
		}
		int rootIdx = -1;
		StringTokenizer st = new StringTokenizer(reader.readLine());
		for (int i = 0; i < n; i++) {
			int parent = Integer.parseInt(st.nextToken());
			if(parent == -1) {
				rootIdx = i;
				continue;
			}
			nodes[parent].children.add(i);
		}
		removeNode = reader.readInt();
		if(removeNode == rootIdx) {
			System.out.println(0); return;
		}
		remove(rootIdx);
		System.out.println(countLeaf(rootIdx));
	}

	private static void remove(int index) {
		Node here = nodes[index];
		for (int i = 0; i < here.children.size(); i++) {
			int child = here.children.get(i);
			if(child == removeNode) {
				here.children.remove(i); return;
			}
			remove(child);
		}
	}

	private static int countLeaf(int index) {
		Node here = nodes[index];

		if(here.children.size() == 0) return 1;

		int cnt = 0;
		for (int i = 0; i < here.children.size(); i++) {
			int child = here.children.get(i);
			cnt += countLeaf(child);
		}
		return cnt;
	}
}
class Node {
	int id;
	List<Integer> children = new ArrayList<>();
	public Node(int id) {
		this.id = id;
	}
}
class InputReader {
	private BufferedReader br;

	public InputReader() {
		br = new BufferedReader(new InputStreamReader(System.in));
	}

	public InputReader(String filepath) {
		try {
			br = new BufferedReader(new FileReader(filepath));
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		}
	}

	public List<Character> readLineIntoCharList() throws IOException {
		List<Character> l = new ArrayList<>();
		while(true) {
			int readVal = br.read();
			if(readVal == '\n' || readVal == -1) break;
			l.add((char)readVal);
		}
		return l;
	}

	public boolean ready() throws IOException {
		return br.ready();
	}

	public String readLine() throws IOException {
		return br.readLine();
	}
	public int readInt() throws IOException {
		return Integer.parseInt(readLine());
	}
	public Long readLong() throws IOException {
		return Long.parseLong(readLine());
	}
}