🐨이진 검색 트리

이진 검색 트리

이 문제는 배열 P를 사용하여 이진 검색 트리를 생성한 후, 모든 노드의 높이 합을 구하는 알고리즘을 구현하는 문제입니다. 문제 원문은 아래와 같습니다:

문제 원문

시간 제한: 2 초
메모리 제한: 256 MB
문제:
P는 크기가 N인 배열이다. P에는 0보다 크거나 같고, N-1보다 작거나 같은 정수가 중복 없이 채워져 있다. 이진 검색 트리는 루트가 있는 이진 트리로, 각각의 노드에 정수 값이 저장되어 있는 트리이다. 이진 검색 트리를 P배열을 이용해서 만드는 법은 다음과 같다. 일단 root를 만들고 거기에 P[0]의 값을 넣은 후에 다음과 같은 과정을 거친다.

for (int i=1; i<=n-1; i++) {
    insert(root, P[i]);
}

void insert(Vertex V, int X) {
    if (x < V에 저장되어 있는 수) {
        if (V가 왼쪽 자식이 있으면) {
            insert(V의 왼쪽 자식, X);
        } else {
            V의 왼쪽 자식을 새로 만들고, 그 곳에 X를 저장함
        }
    } else {
        if (V가 오른쪽 자식이 있으면) {
            insert(V의 오른쪽 자식, X);
        } else {
            V의 오른쪽 자식을 새로 만들고, 그 곳에 X를 저장함
        }
    }
}

N과 배열 P에 있는 수가 주어졌을 때, P로 이진 검색 트리를 만들었을 때 모든 노드의 높이의 합을 출력하는 프로그램을 작성하시오. 트리의 높이는 루트에서부터의 거리 + 1이다.

예제 입력/출력

입력1
10
9
1
4
3
2
5
6
7
8
0

출력1
40

입력2
10
6
3
2
7
9
4
8
1
0
5

출력2
31

입력3
10
0
1
2
3
4
5
6
7
8
9

출력3
55

정답 코드

import bisect
import array

def main():
    n,*nodes = map(int,open(0).read().split())
    h = array.array('i', [0] * n)
    bst_list = array.array('i')
    total_height_sum = 0
    
    for node in nodes:
        height = get_height(node, h, bst_list)
        total_height_sum += height

    print(total_height_sum)

def get_height(node, h, bst_list):
    idx = bisect.bisect_left(bst_list, node)

    left_height = h[bst_list[idx - 1]] if idx > 0 else 0
    right_height = h[bst_list[idx]] if idx < len(bst_list) else 0

    h[node] = max(left_height, right_height) + 1
    bst_list.insert(idx, node)

    return h[node]

if __name__ == "__main__":
    main()

코드 설명

이 코드는 이진 검색 트리에서 노드의 높이를 계산하는 효율적인 방법을 제공합니다. 아래는 주요 알고리즘과 관련 개념입니다:

1. `bisect` 모듈 사용

Python의 bisect 모듈은 정렬된 리스트에서 삽입할 위치를 효율적으로 찾는 도구입니다. 여기서는 bisect_left를 사용해 노드의 올바른 삽입 위치를 찾습니다. 시간 복잡도는 O(log N)입니다.

idx = bisect.bisect_left(bst_list, node)

2. 노드 높이 계산

현재 노드의 높이는 좌우 자식 중 최대 높이에 1을 더한 값입니다. 이를 통해 트리에서 노드의 위치에 따른 높이를 계산합니다.

h[node] = max(left_height, right_height) + 1

3. 이진 검색 트리 갱신

정렬된 상태를 유지하기 위해 bst_list에 새 노드를 적절한 위치에 삽입합니다.

bst_list.insert(idx, node)

결론

이 문제는 이진 검색 트리의 특성과 Python의 효율적인 정렬 알고리즘 도구를 활용하여 주어진 제약 조건 내에서 높이의 합을 계산할 수 있도록 설계되었습니다. 정렬된 배열을 유지하며 높이를 계산하는 과정은 실무에서도 효율적인 데이터 처리에 유용한 접근 방식입니다.

저작자표시 비영리 변경금지

붉은 각설탕

🐨이진 검색 트리