[Python] 31530번 새로운 AVL 트리 만들기

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

31530번: 새로운 AVL 트리 만들기

---

 

24/03/21

 

 

간단한 DP문제로, 토글링의 방법으로 해결할 수 있는 문제다.

---

 

문제 접근 방식:

 

 

DP테이블을 다음과 같이 정의하자.

 

$$DP[i][j] = \textrm{높이가 }i\textrm{이고 가능한 균형값의 상태가 }j\textrm{일 때의 경우의 수}$$

 

이때, 가능한 균형값의 상태는 다음과 같이 정의한다.

 

$$j = 0, 1, 2 \rightarrow \{-1\} , \{0\} , \{1\}$$

$$j = 3, 4, 5 \rightarrow \{-1, 0\} , \{-1, 1\}, \{0, 1\}$$

$$j = 6 \rightarrow \{ -1, 0, 1\}$$

 

우리는 높이가 $h$이고, 가능한 균형값의 상태가 $j$인 트리가 주어질 때, 이 트리를 확장하여 $h+1$의 높이를 가지는 트리를 구성하고자 한다.

 

$j = 0$일 때, 즉, 가능한 균형값의 상태가 $\{-1\}$라면 트리를 확장하려고 할 때, 루트 노드 왼쪽에만 기존의 트리를 붙일 수 있으므로, 점화식은 다음과 같이 정의된다.

 

$$DP[i][0] = DP[i-1][0]$$

 

이는 $j=1, 2$일 때도 비슷하다. 이때는 루트 노드 양쪽 모두에 붙이는 경우와 루트 노드 오른쪽에 붙이는 경우만 존재하므로 점화식이 다음과 같이 정의된다.

 

$$DP[i][1] = DP[i-1][1] ; DP[i][2] = DP[i-1][2]$$

 

$j=3$인 경우, 루트 노드가 가능한 균형값은 $\{-1, 0\}$으로 2개이다.

 

루트 노드의 균형값이 $0$이라면 왼쪽 서브트리와 오른쪽 서브트리의 높이가 $i-1$로 같고, 서브트리에 올 수 있는 트리 개수는 $DP[i-1][3]$이다.

 

루트 노드의 균형값이 $-1$이라면 왼쪽 서브트리는 높이가 $i-1$이고 오른쪽 서브트리는 높이가 $i-2$이며, 각 서브트리에 올 수 있는 트리 개수는 $DP[i-2][3]$과 $DP[i-1][3]$이다.

 

따라서, 점화식은 $DP[i][3] = DP[i-2][3]*DP[i-1][3] + DP[i-1][3]*DP[i-1][3]$로 정의된다.

 

비슷하게, $j = 4, 5, 6$일 때도 점화식을 정의할 수 있으며, 이를 통해 모든 AVL 트리의 경우의 수를 구할 수 있다.

 

구할 때마다 소수로 모듈러 연산을 취해주어야 함에 유의하자.

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 더 보기를 누르면 확인할 수 있다.

# 31530번 새로운 AVL 트리 만들기
# DP
'''
접근 방법:
DP[i][j] = 높이 i이고 가능한 균형값 상태가 j일때
j = 0, 1, 2 -> {-1}, {0}, {1}
j = 3, 4, 5 -> {-1,0}, {-1,1}, {0,1}
j = 6 -> {-1, 0, 1}
DP[i][0] = DP[i-1][0]
DP[i][1] = DP[i-1][1]
DP[i][2] = DP[i-1][2]
DP[i][3] = DP[i-2][3]*DP[i-1][3] + DP[i-1][3]*DP[i-1][3]
DP[i][4] = DP[i-2][4]*DP[i-1][4] + DP[i-1][4]*DP[i-2][4]
DP[i][5] = DP[i-2][5]*DP[i-1][5] + DP[i-1][5]*DP[i-1][5]
DP[i][6] = DP[i-2][6]*DP[i-1][6] + DP[i-1][6]*DP[i-1][6] + DP[i-1][6]*DP[i-2][6]
'''
import sys
input = sys.stdin.readline

MOD = 1_000_000_000 + 7
DP = [[0 for _ in range(7)] for _ in range(1_000_000+1)]
DP[0] = [1, 1, 1, 1, 1, 1, 1]
DP[1] = [1, 1, 1, 2, 2, 2, 3]
for i in range(2, 1_000_000+1):
    DP[i][0] = DP[i-1][0] % MOD
    DP[i][1] = DP[i-1][1] % MOD
    DP[i][2] = DP[i-1][2] % MOD
    DP[i][3] = (DP[i-2][3]*DP[i-1][3] + DP[i-1][3]*DP[i-1][3]) % MOD
    DP[i][4] = (DP[i-2][4]*DP[i-1][4] + DP[i-1][4]*DP[i-2][4]) % MOD
    DP[i][5] = (DP[i-2][5]*DP[i-1][5] + DP[i-1][5]*DP[i-1][5]) % MOD
    DP[i][6] = (DP[i-2][6]*DP[i-1][6] + DP[i-1][6]*DP[i-1][6] + DP[i-1][6]*DP[i-2][6]) % MOD
ans = []
T = int(input())
for _ in range(T):
    h, S = map(int, input().split())
    K = input().rstrip()
    if K == '-1':
        ans.append(DP[h-1][0])
    elif K == '0':
        ans.append(DP[h-1][1])
    elif K == '1':
        ans.append(DP[h-1][2])
    elif K == '-1 0':
        ans.append(DP[h-1][3])
    elif K == '-1 1':
        ans.append(DP[h-1][4])
    elif K == '0 1':
        ans.append(DP[h-1][5])
    else:
        ans.append(DP[h-1][6])
print(*ans, sep='\n')