균형이진탐색트리(AVL, Balanced BST)는 모든 노드에 대해서 노드의 왼쪽 sub tree와 오른쪽 sub tree의 높이차가 1 이하인 BST를 뜻한다.
class Node:
def __init__(self, key = None, parent = None, left = None, right = None):
self.key = key
self.parent = parent
self.left = left
self.right = right
self.height = 0
def __str__(self):
return str(self.key)
class BST:
def __init__(self):
self.root = None
self.size = 0
self.height = 0
def __len__(self):
return self.height
def __iter__(self):
return self.root.__iter__()
def __str__(self):
return " - ".join(str(k) for k in self)
def preorder(self, v): #MLR
if v:
print(v.key, end=' ')
self.preorder(v.left)
self.preorder(v.right)
def inorder(self, v): #LMR
if v:
self.inorder(v.left)
print(v.key, end= ' ')
self.inorder(v.right)
def postorder(self, v): #LRM
if v:
self.postorder(v.left)
self.postorder(v.right)
print(v.key, end=' ')
def find_loc(self, key): # key가 트리에 있으면 해당 노드 리턴, 없으면 해당 노드가 삽입될 곳의 부모노드 리턴
if self.size == 0: # 트리가 비었을때 None 리턴
return None
p = None
v = self.root
while v:
if v.key == key:
return v
elif v.key > key:
p = v
v = v.left
else:
p = v
v = v.right
return p # 삽입될 곳의 부모노드 p 리턴
def search(self, key):
p = self.find_loc(key)
if p and p.key == key: # p가 있고, key값이 같을 때 p 리턴
return p
else:
return None # 없으면 None 리턴
def insert(self, key):
v = Node(key) # key값을 가진 노드 v 생성
if self.size == 0: # 빈 트리이면 v는 루트노드
self.root = v
else:
p = self.find_loc(key)
if p and p.key != key: # p가 존재하지만 key값이 다를 때 v의 위치 정하기
if p.key > key:
p.left = v
else:
p.right = v
v.parent = p # p,v는 부모자식 관계 설정
self.size += 1
return v #삽입된 노드 v 리턴
def update_node_height(self,v):
if v:
l = v.left.height if v.left else -1
r = v.right.height if v.right else -1
v.height = max(l, r) + 1
def update_height(self, v):
while v != None:
self.update_node_height(v)
v = v.parent
def deleteByMerging(self, x):
if x == None:
return None
a = x.left
b = x.right
pt = x.parent
# c는 삭제될 x 위치에 올 노드
# s는 균형이 깨질 가능성이 있는 첫 번째 노드
if a == None: # x의 왼쪽 부트리가 없을 때
c = b # x의 오른쪽 부트리인 b를 그대로 c위치로 이동
s = pt # x가 삭제되었기때문에 균형이 깨질 가능성이 큰 노드는 x의 부모노드인 pt
else: # a가 있을 때
c = a # a를 c자리에 올림
m = a # m을 a에서 가장 큰 노드로 두고 그걸 찾기 위함.
while m.right:
m = m.right
m.right = b # m의 오른쪽 자식으로 b를 연결
if b: # b가 있으면 b의 부모를 m으로 연결
b.parent = m
s = m # m에 b가 추가되었으므로 m은 균형이 깨질 수 있다.
if self.root == x: # x가 루트노드이면 c를 루트노드로 올림
if c:
c.parent = None
self.root = c
else: # x가 루트노드가 아닐 때 x자리에 c노드 넣기
if pt.left == x:
pt.left = c
else:
pt.right = c
if c: # c가 있으면 c의 부모를 pt로 연결
c.parent = pt
self.size -= 1
return s
def deleteByCopying(self, x):
a = x.left
b = x.right
pt = x.parent
if a: #x의 왼쪽 부트리만 있는 경우
y = a
while y.right: #y를 a에서 가장 큰 값으로 설정
y = y.right
x.key = y.key #값 Copy
if y.left: #y의 왼쪽 부트리가 있는 경우
if y == y.parent.right: #y가 y.parent의 오른쪽에 있을 때
y.parent.right = y.left
else: #y가 y.parent의 왼쪽에 있을 때 *****y가 x.left일때 즉, while문을 한번도 안돌았을 경우우
y.parent.left = y.left
y.left.parent = y.parent
else: #y가 리프노드일 경우 y와 y의 부모자식 관계를 없앰
if y == y.parent.right:
y.parent.right = None
else:
y.parent.left = None
self.size -= 1
elif a == None and b != None: #왼쪽은 없고 오른쪽만 있는 경우
y = b
while y.left: #y를 b에서 가장 작은 값으로 설정
y = y.left
x.key = y.key #값 Copy
if y.right: # x의 오른쪽 부트리에서 가장 작은 값 y에 오른쪽 자식이 있는 경우
if y == y.parent.left: #y가 y의 부모의 왼쪽에 있는 경우
y.parent.left = y.right
else:
y.parent.right = y.right #y가 y의 부모의 오른쪽에 있는 경우
y.right.parent = y.parent
else: #y가 리프노드일 경우 y와 y의 부모자식 관계를 없앰
if y == y.parent.left:
y.parent.left = None
else:
y.parent.right = None
self.size -= 1
elif a == None and b == None: #x가 리프노드일 경우
if pt: #x의 부모노드가 있는 경우
if x == pt.right: #x가 x의 부모의 오른쪽 자식일 때 관계 지우기
pt.right = None
elif x == pt.left: #x가 x의 부모의 왼쪽 자식일 때 관계 지우기
pt.left = None
else: #x가 리프노드인데 부모노드가 없는 경우 = x는 루트노드
self.root = None
self.size -= 1
return pt
def rotateRight(self, z):
if z == None: # 돌릴게 없다.
return
x = z.left # z의 왼쪽 자식 x
if x == None: # x가 없으면 돌릴게 없다.
return
b = x.right # x의 오른쪽 자식 b
x.parent = z.parent # 회전 시작. x의 부모가 z의 부모노드가 되도록.
if z.parent: # z의 부모가 있을때 z의 자리에 x를 넣음.
if z.parent.left == z:
z.parent.left = x
else:
z.parent.right = x
x.right = z # x의 오른쪽에 z
z.parent = x # x,z 부모자식관계 형성
z.left = b #x의 오른쪽 자식이였던 b가 이젠 z의 왼쪽 자식으로
if b: # b,z 부모자식관계 형성
b.parent = z
if z == self.root: # z가 루트노드였을 경우
self.root = x
self.update_node_height(x) # 높이 업데이트
self.update_node_height(z)
def rotateLeft(self, x):
if x == None:
return
z = x.right # z는 x의 오른쪽 자식
if z == None:
return
b = z.left # b는 z의 왼쪽 자식
z.parent = x.parent # 회전 시작. z의 부모노드가 x의 부모노드가 되도록.
if x.parent: # x의 부모가 있을 때 x의 자리에 z를 넣음.
if x.parent.left == x:
x.parent.left = z
else:
x.parent.right = z
z.left = x # z의 왼쪽에 z
x.parent = z # x,z 부모자식관계 형성성
x.right = b # z의 왼쪽 자식이였던 b가 이젠 x의 오른쪽 자식으로
if b: # b,x 부모자식관계 형성
b.parent = x
if x == self.root: # x가 루트노드였을 경우
self.root = z
self.update_node_height(x) # 높이 업데이트
self.update_node_height(z)
class AVL(BST):
def __init__(self):
super().__init__()
def rebalance(self, x, y, z): # rebalancing한 후 원래 z 노드 자리로 올라온 노드 리턴
if z == None: # z가 없으면 아무것도 안함.
return
if y == z.left and x == y.left: #왼쪽으로 z-y-x 일직선
self.rotateRight(z)
return y
elif y == z.left and x == y.right: #왼쪽으로 z-y-x 삼각형
self.rotateLeft(y)
self.rotateRight(z)
return x
elif y == z.right and x == y.right: #오른쪽으로 z-y-x 일직선
self.rotateLeft(z)
return y
elif y == z.right and x == y.left: #오른쪽으로 z-y-x 삼각형
self.rotateRight(y)
self.rotateLeft(z)
return x
def insert(self, key):
v = super(AVL, self).insert(key) #BST의 insert 호출
current = v # v를 current라고 저장 (return v 하기 위함.)
while current: # current가 있을 때
left_height = current.left.height if current.left else -1 # current의 오른쪽 왼쪽 부트리 높이 구하기
right_height = current.right.height if current.right else -1
if abs(left_height - right_height) > 1: #균형이 깨졌는지 체크
z = current #균형이 깨진 첫 번째 노드 z로 지정
if left_height>right_height: # 더 노드가 많은쪽에 z의 자식으로 y 지정
y = z.left
else:
y = z.right
if y: # y가 있을 때
y_left_height = y.left.height if y.left else -1 # y의 오른쪽 왼쪽 부트리 높이 구하기
y_right_height = y.right.height if y.right else -1
if y_left_height>y_right_height: # 더 노드가 많은쪽에 y의 자식으로 x 지정
x = y.left
else:
x= y.right
if x and y and z: # x, y, z가 다 있으면 균형맞추기 시작
self.rebalance(x,y,z)
break # 균형 다 맞췄으면 중단
current = current.parent # 삽입된 노드의 부모로 균형이 깨질수 있는 노드 추적하기
return v # 삽입된 노드 리턴
def delete(self, u):
v = super(AVL, self).deleteByCopying(u) #BST의 deleteByCopying 호출
while v: # v가 있는 경우
v_left_height = v.left.height if v.left else -1
v_right_height = v.right.height if v.right else -1
if abs(v_left_height - v_right_height) > 1: # v에서 균형이 께졌는지 확인
z = v # 균형이 깨진 첫 번째 노드 z로 지정
if z.left.height >= z.right.height: # 더 노드가 많은 쪽에 z의 자식으로 y 지정
y = z.left
else:
y = z.right
if y: # y가 있을 때
y_left_height = y.left.height if y.left else -1
y_right_height = y.right.height if y.right else -1
if y_left_height >= y_right_height: # 더 노드가 많은 쪽에 y의 자식으로 x 지정
x = y.left
else:
x = y.right
if x and y and z: # x, y, z가 다 있으면 균형맞추기 시작작
new_top = self.rebalance(x,y,z)
v = new_top.parent #균형 다 맞춰진 노드의 부모를 v로 재 지정
continue # while문으로 돌아가서 부모노드를 추적하면서 rebalancing하기
w = v # 마지막으로 검사한 노드 기억하기
v = v.parent # 루트노드까지 균형 재검사
self.root = w # 회전으로 루트노드가 바뀌었을 때 루트노드 정확히 설정하기AVL 트리 수행시간, 회전 횟수 정리
해적왕이 될 사나이