Skip to content

Instantly share code, notes, and snippets.

@nuovothoth

nuovothoth/AVLTree.kt

Last active July 22, 2024 12:49
Show Gist options
  • Save nuovothoth/81b47c2f7198460a9fb25beefc4a89e5 to your computer and use it in GitHub Desktop.
Save nuovothoth/81b47c2f7198460a9fb25beefc4a89e5 to your computer and use it in GitHub Desktop.
Download ZIP
AVLTree
Raw
AVLTree.kt
package datastructure
import java.util.*
class AVLTree {
var root: Node? = null
private set
fun rebalance(node: Node): Node {
val bf = getBF(node) //BF 계산
var returnNode = node
when {
bf > 1 -> { // 왼쪽 서브 노드의 높이로 인해 불균형 발생
if ((node.left?.left?.height?.plus(1) ?: 0) > (node.left?.right?.height?.plus(1) ?: 0)) { // LL
returnNode = rightRotate(node)
} else { // LR
node.left = node.left?.let { leftRotate(it) }
returnNode = rightRotate(node)
}
}
bf < -1 -> { // 오른쪽 서브 노드의 높이로 인해 불균형 발생
if ((node.right?.right?.height?.plus(1) ?: 0) > (node.right?.left?.height?.plus(1) ?: 0)) { // RR
returnNode = leftRotate(node)
} else { // RL
node.right = node.right?.let {
rightRotate(it)
}
returnNode = leftRotate(node)
}
}
else -> node
}
return returnNode
}
fun rightRotate(node: Node): Node {
val firstNode = node
val secondNode = node.left
val rightSubNodeOfSecondNode = secondNode!!.right
secondNode.right = firstNode
firstNode.left = rightSubNodeOfSecondNode
firstNode.height = calcHeight(firstNode)
secondNode.height = calcHeight(secondNode)
return secondNode
}
fun leftRotate(node: Node): Node {
val firstNode = node
val secondNode = node.right
val leftSubNodeOfSecondNode = secondNode!!.left
secondNode.left = firstNode
firstNode.right = leftSubNodeOfSecondNode
firstNode.height = calcHeight(firstNode)
secondNode.height = calcHeight(secondNode)
return secondNode
}
fun insert(key: Int, node: Node?): Node? {
if (root == null) { // 루트 노드조차 생성되어 있지 않으면 생성
root = Node(key = key)
return root
}
if (node == null) return Node(key = key) // 해당 노드 자리에 아무 노드도 없으면 여기에 노드 생성해서 삽입(하도록 리턴)
when {
node.key > key -> node.left = insert(key, node.left)
node.key < key -> node.right = insert(key, node.right)
else -> throw IllegalArgumentException() // AVL 노드는 노드의 값이 유니크해야함
}
node.height = calcHeight(node) // 삽입 후 높이 재측정
val returnNode = rebalance(node) // 리밸런싱 진행 (필요 시)
// node가 루트였는데 리밸런싱이 진행되었으면 root가 가리키고 있는 것이 더이상 루트가 아닐 수 있음
if (findParentNode(node.key) == null) {
root = returnNode
}
return returnNode
}
fun delete(key: Int, node: Node?): Node? {
if (node == null) return null
var replacedNode = node
when {
node.key > key -> node.left = delete(key, node.left)
node.key < key -> node.right = delete(key, node.right)
else -> {
when {
node.left != null -> { // 삭제할 노드의 왼쪽 서브노드 존재
val leftNode = node.left!!
val maxNode = findMaxKeyNode(leftNode) // 왼쪽 서브노드의 가장 오른쪽 노드는 서브노드 중 가장 큰 값을 가짐
val maxNodeParentNode = findParentNode(maxNode.key)!! // maxNode의 parentNode는 나중에 maxNode와의 관계를 끊어야 함
maxNode.right = node.right // node 오른쪽 서브노드가 있다면 node를 대체할 maxNode의 오른쪽에 해당 서브노드를 연결
connectToParentNodeOfNode(maxNode, node) // 삭제할 노드의 부모 노드의 새로운 자식 노드로 maxNode를 연결
// node의 왼쪽 노드는 maxNode보다 항상 작기 때문에 maxNode의 왼쪽에 위치
// 그러나 node의 왼쪽 서브노드가 maxNode 하나뿐이였다면 사이클을 만들지 말아야 함
if (maxNode != leftNode) {
maxNode.left = leftNode
disconnectFromParent(maxNode, maxNodeParentNode)
}
// node를 대체하여 반환될 새로운 노드
replacedNode = maxNode
}
node.right != null -> { // 삭제할 노드의 오른쪽 서브노드 존재
val rightNode = node.right!!
val minNode = findMinKeyNode(rightNode) // 오른쪽 서브노드의 가장 왼쪽 노드는 서브노드 중 가장 작은 값을 가짐
val minNodeParentNode = findParentNode(minNode.key)!! // minNode의 parentNode는 나중에 minNode와의 관계를 끊어야 함
minNode.left = node.left // node 왼쪽 서브노드가 있다면 node를 대체할 minNode의 왼쪽에 해당 서브노드를 연결
connectToParentNodeOfNode(minNode, node) // 삭제할 노드의 부모 노드의 새로운 자식 노드로 minNode를 연결
// node의 오른쪽 노드는 minNode보다 항상 크기 때문에 maxNode의 오른쪽에 위치
// 그러나 node의 오른쪽 서브노드가 minNode 하나뿐이였다면 사이클을 만들지 말아야 함
if(minNode != rightNode) {
minNode.right = rightNode
disconnectFromParent(minNode, minNodeParentNode)
}
// node를 대체하여 반환될 새로운 노드
replacedNode = minNode
}
else -> { // 삭제할 노드가 리프노드인 경우
return null
}
}
}
}
replacedNode.height = calcHeight(replacedNode) // 삭제 후 높이 재측정
val returnNode = rebalance(replacedNode) // 리밸런싱 진행 (필요 시)
// returnNode가 새로운 루트노드가 되었을 수 있음
if (findParentNode(returnNode.key) == null) {
root = returnNode
}
return returnNode
}
fun connectToParentNodeOfNode(node: Node, targetNode: Node) {
val parentNode = findParentNode(targetNode.key)
if ((parentNode?.left?.key ?: -1) == targetNode.key) {
parentNode?.left = node
} else if ((parentNode?.right?.key ?: -1) == targetNode.key) {
parentNode?.right = node
}
}
//한 노드의 BF: 왼쪽 서브트리의 높이 - 오른쪽 서브트리의 높이
fun getBF(node: Node): Int {
val leftChildHeight = node.left?.height?.plus(1) ?: 0
val rightChildHeight = node.right?.height?.plus(1) ?: 0
return leftChildHeight-rightChildHeight
}
fun calcHeight(node: Node): Int {
return Math.max(node.left?.height?.plus(1) ?: 0, node.right?.height?.plus(1) ?: 0)
}
fun disconnectFromParent(targetNode: Node, targetParentNode: Node) {
when(targetNode.key) {
targetParentNode.left?.key -> {
targetParentNode.left = null
}
else -> {
targetParentNode.right = null
}
}
}
fun findMaxKeyNode(node: Node): Node {
return node.right?.let {
findMaxKeyNode(it)
} ?: run {
node
}
}
fun findMinKeyNode(node: Node): Node {
return node.left?.let {
findMinKeyNode(it)
} ?: run {
node
}
}
fun findParentNode(key: Int): Node? {
var currentNode = root
while(true) {
if (currentNode == null) return null
if (currentNode.key > key) {
when (currentNode.left?.key) {
null -> return null
key -> return currentNode
else -> currentNode = currentNode.left
}
} else if (currentNode.key < key) {
when (currentNode.right?.key) {
null -> return null
key -> return currentNode
else -> currentNode = currentNode.right
}
} else {
return null
}
}
}
fun bfs(node: Node?) {
if (node == null) return
val queue: Queue<Pair<Node, Int>> = LinkedList()
queue.offer(node to 0)
var currentLevel = node.height
while(queue.isNotEmpty()) {
val pair = queue.poll()
if (pair.second != currentLevel) {
currentLevel = pair.second
println("")
}
print("${pair.first.key}(${pair.first.height}) ")
pair.first.left?.let {
queue.offer(it to currentLevel+1)
}
pair.first.right?.let {
queue.offer(it to currentLevel+1)
}
}
}
}
fun main(args: Array<String>) {
val avlTree = AVLTree()
avlTree.insert(3, avlTree.root)
avlTree.insert(2, avlTree.root)
avlTree.insert(1, avlTree.root)
avlTree.insert(7, avlTree.root)
avlTree.insert(6, avlTree.root)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.insert(5, avlTree.root)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.insert(4, avlTree.root)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.delete(7, avlTree.root)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.delete(5, avlTree.root)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.delete(4, avlTree.root)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.delete(6, avlTree.root)
avlTree.bfs(avlTree.root)
println("\n-----------------")
}
/*
2(2)
1(0) 6(1)
3(0) 7(0)
-----------------
3(2)
2(1) 6(1)
1(0) 5(0) 7(0)
-----------------
3(3)
2(1) 6(2)
1(0) 5(1) 7(0)
4(0)
-----------------
3(2)
2(1) 5(1)
1(0) 4(0) 6(0)
-----------------
3(2)
2(1) 4(1)
1(0) 6(0)
-----------------
3(2)
2(1) 6(0)
1(0)
-----------------
2(1)
1(0) 3(0)
-----------------
*/
Raw
Node.kt
data class Node(
var height: Int = 0, // 자신을 루트노드로 가정할 때의 높이
var left: Node? = null,
var right: Node? = null,
val key: Int
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment