[백준] 1068 트리, python, C++

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

1068번: 트리

import sys; input = sys.stdin.readline

N = int(input())
node_parent = [-1 for _ in range(N)]
node_son = [[] for _ in range(N)]
root = 0
for son, parent in enumerate(map(int, input().split())):
    if parent == -1:
        root = son
        continue
    node_son[parent].append(son)
removed_node = int(input())
answer = 0
def preorder(cur_node):
    if cur_node == removed_node:
        return
    if not len(node_son[cur_node]) or (len(node_son[cur_node]) == 1 and node_son[cur_node][0] == removed_node):
        global answer
        answer += 1
        return
    for son in node_son[cur_node]:
        preorder(son)

preorder(root)
print(answer)

오랜만에 푸는 트리 문제이다.

전위 순회하는 방식으로 문제를 풀었는데 제거된 노드는 스킵할수 있도록 최상단에 조건문을 걸어두었다.

이런 방식은 한가지 문제점이 있는데 리프노드를 삭제했을때 그 부모노드가 다시 리프노드가 되는 경우 정답을 찾지 못한다.(ex. 2 / -1 0 / 1) 이 경우는 조건문으로 한번더 처리를 해서 통과하였다.

#include<iostream>
#include<vector>
using namespace std;

int answer = 0;

void preorder(int cur_node, int removed_node, vector<vector<int>> node_son) {
	if (cur_node == removed_node) return ;
	if ((node_son[cur_node].empty())
		|| ((1 == node_son[cur_node].size()) && removed_node == node_son[cur_node][0])
	) {
		answer += 1;
		return;
	}
	for (int i = 0; i < node_son[cur_node].size(); i++)
	{
		int son = node_son[cur_node][i];
		preorder(son, removed_node, node_son);
	}
}

int main() {
	// 0 입력
	int N, root, removed_node;
	cin >> N;
	vector<int> node_parent(N);
	vector<vector<int>> node_son(N);
	
	for (int i = 0; i < N; i++)
	{
		cin >> node_parent[i];
		if (node_parent[i] == -1) {
			root = i;
			continue;
		}
		// add son
		node_son[node_parent[i]].push_back(i);
	}
	cin >> removed_node;

	// 1 탐색
	preorder(root, removed_node, node_son);
	
	// 2 출력
	cout << answer;

	return 0;
}

알고리즘은 똑같은데 std에 익숙해지기 위해서 만들었다