[알고리즘] 동적 계획법 - 최적 이진 탐색 트리 코드 (DP - Optimal BST Code)

문제)

Description

교재와 강의자료를 참고하여 Algorithm 3.9/3.10의 구현을 완성하시오.

원소의 개수 n, 키의 값 K, 원소의 탐색 빈도값의 배열 p가 주어질 때

A, R 행렬의 값을 구해서 출력하고,

R 행렬을 이용하여 구축할 수 있는 이진탐색트리의

preorder, inorder 순회 탐색 결과를 출력하시오.

Input

첫 번째 줄에 key의 개수 n이 주어진다.

두 번째 줄에 n 개의 key 값이 주어진다. (key 값은 정렬되어 있다고 가정해도 된다.)

세 번째 줄에 n 개의 빈도값 p가 주어진다. (p[i] 값은 양의 정수값으로 주어진다.)

Output

먼저 행렬 A의 윗 부분 삼각형을 출력한다. (0을 포함)

다음으로 행렬 R의 윗 부분 삼각형을 출력한다. (0을 포함)

A와 R을 출력한 후에 최적 이진탐색트리에서 평균검색시간의 최적값을 출력한다.

다음 줄에 최적 이진탐색트리의 preorder 순회 탐색 결과를 출력한다.

다음 줄에 최적 이진탐색트릴의 inorder 순회 탐색 결과를 출력한다.

Sample Input 1

4
10 20 30 40
3 3 1 1

Sample Output 1

0 3 9 11 14
0 3 5 8
0 1 3
0 1
0
0 1 1 2 2
0 2 2 2
0 3 3
0 4
0
14
20 10 30 40
10 20 30 40

코드)

import sys
input = sys.stdin.readline


class node:
    # 트리를 구성하는 노드
    def __init__(self, key):
        self.left = None
        self.right = None
        self.key = key


def minimum(i ,j):
    # i부터 j까지의 노드 중에 최적의 트리를 찾아줌
    minvalue = 999999
    minidx = 0
    for k in range(i, j + 1):
        # k는 i가 루트일때부터, j가 루트일때까지 계산한다.
        value = A[i][k - 1] + A[k + 1][j] + sum(p[i : j + 1])
        if minvalue > value:
            minvalue = value
            minidx = k

    return minvalue, minidx
        


def func(n):
    # 최적 값을 찾는 함수
    for diagonal in range(1, n):        # 대각선은 총 n개가 있음
        for i in range(1, n - diagonal + 1):
            # diagonal = 1 일때, i는 1부터 n - 1까지 가야함
            # diagonal = 2 일때, i는 1부터 n - 2까지 가야함
            j = i + diagonal        # diagonal = 1 일때, j는 i의 바로 다음 인덱스임
            A[i][j], R[i][j] = minimum(i, j)


def create_tree(start, end):
    k = R[start][end]
    if k == 0:
        return None
    else:
        root = node(key[k])
        root.left = create_tree(start, k - 1)
        root.right = create_tree(k + 1, end)

        return root


def inorder(root):
    if root != None:
        inorder(root.left)
        print(root.key, end=' ')
        inorder(root.right)


def preorder(root):
    if root != None:
        print(root.key, end=' ')
        preorder(root.left) 
        preorder(root.right)            

        
n = int(input())
key = [0] + list(map(int, input().split()))       # key 값을 입력 받음
p = [0] + list(map(int, input().split()))     # 확률 값을 입력 받음

A = [[0] * (n+1) for _ in range(n+2)]       # 최적 값을 저장하는 배열
R = [[0] * (n+1) for _ in range(n+2)]       # 최적 값이 되는 루트를 저장하는 배열
# 이때 행은 1부터 n+1이고, 열은 1부터 n까지인 이유는, k에 값이 들어갈 때 행은 초과할 수 있지만, 열은 초과할 수 없기 때문이다.
# 예를 들어 n이 최고 루트라면 A[1][n-1]까지의 값과 A[n+1][n]의 값을 더해야 함. 따라서 행을 초과할 수 있음
# 하지만 1이 최고 루트라면 A[1][0] + A[2][n]인데, 이때 열은 (아래로) 초과할 수 없음

for i in range(1, n + 1):
    A[i][i] = p[i]
    R[i][i] = i

func(n)

for i in range(1, n + 2):
    print(*A[i][i - 1:])
for i in range(1, n + 2):
    print(*R[i][i - 1:])
print(A[1][n])

# 트리를 만들어야 함
root = create_tree(1, n)
preorder(root)
print()
inorder(root)
print()