[Python/BOJ] 23844. 트리 정리하기

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

23844번: 트리 정리하기

문제

동빈이가 사는 나라의 나무들은 트리 모양을 하고 있다.

다가오는 크리스마스를 기념해서, 동빈이는 나무를 다듬어서 집에 트리를 장식하기로 했다.

동빈이가 가지고 있는 트리는 N개의 노드들이 N-1개의 간선으로 연결되어 있다. 각 노드들은 1~N까지의 번호가 정해져 있고, 루트 노드의 번호는 1이다.

트리에서 루트 노드의 레벨은 0이고 아래로 갈수록 1씩 증가한다. 동빈이는 모든 노드에 장식을 달고 싶었지만, 트리의 두께가 K를 넘으면 집에 들어가지 않는다.

이때 트리의 두께란, 같은 레벨에 있는 노드의 개수 중 최댓값이다.

결국, 일부 노드를 제거해 트리의 두께가 K이하가 되도록 만들어야 한다. 트리를 다듬을 때 부모 노드를 제거하면 모든 자식 노드도 제거된다.

되도록 많은 장식을 달고 싶은 동빈이는 최대한 많은 노드를 남기려 한다. 여러분이 동빈이를 도와주자.

입력

첫째 줄에 정점의 개수 N과 레벨 당 남길 수 있는 노드의 개수 K가 주어진다. (1 ≤ N ≤ 10,000, 1 ≤ K ≤ N) 

다음 N-1개의 줄에 걸쳐, 두 정수 ai, bi 가 주어진다. 이는 ai번 노드와 bi번 노드가 간선으로 연결되어 있다는 의미이다. (1 ≤ ai, bi ≤ N, ai ≠ bi) 

출력

최대한 많은 노드를 남기고 트리를 정리했을 때 남은 노드의 개수를 출력하시오.

풀이

Code

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**6)
def solution(n, k) :
    # 1. 자식 노드의 수 추출 함수 설정
    def dfs(now, level) :
        global visited
        # 1-1. 방문 처리
        visited |= (1 << now)
        # 1-2.
        for next in graph[now] :
            # 1-2-1. 방문한 경우 continue
            if visited & (1 << next) : continue
            # 1-2-2. 현재 레벨의 노드 수가 k개 이하일 경우 카운팅
            if level_cnt[level] < k : level_cnt[level] += 1
            # 1-2-3. dFS 실행
            dfs(next, level + 1)

    # 2. 그래프 생성 후 간선 정보 입력
    graph = [[] for _ in range(n+1)]
    for _ in range(n-1) :
        a, b = map(int, input().split())
        graph[a].append(b)
        graph[b].append(a)
    level_cnt = [0 for _ in range(n+1)]
    level_cnt[0] = 1
    # 3. dfs 실행
    dfs(1, 1)
    # 4. 결과값 출력
    print(sum(level_cnt))

if __name__ == "__main__" :
    n, k = map(int, input().split())
    visited = 0
    solution(n, k)