Skip to content

Instantly share code, notes, and snippets.

@nuovothoth
Created July 22, 2024 13:18
Show Gist options
  • Save nuovothoth/533e3282904eecff161e2d64c338d1c6 to your computer and use it in GitHub Desktop.
Save nuovothoth/533e3282904eecff161e2d64c338d1c6 to your computer and use it in GitHub Desktop.
AVLTree by chat-GPT
package datastructure
import java.util.*
class AVLNode(var key: Int, var left: AVLNode? = null, var right: AVLNode? = null) {
var height: Int = 1
fun updateHeight() {
height = 1 + maxOf(left?.height ?: 0, right?.height ?: 0)
}
fun balanceFactor(): Int {
return (left?.height ?: 0) - (right?.height ?: 0)
}
}
class AVLTree2 {
var root: AVLNode? = null
private set
fun insert(key: Int) {
root = insertRecursive(root, key)
}
private fun insertRecursive(node: AVLNode?, key: Int): AVLNode? {
if (node == null) {
return AVLNode(key)
}
if (key < node.key) {
node.left = insertRecursive(node.left, key)
} else {
node.right = insertRecursive(node.right, key)
}
node.updateHeight()
return balance(node)
}
fun delete(key: Int) {
root = deleteRecursive(root, key)
}
private fun deleteRecursive(node: AVLNode?, key: Int): AVLNode? {
if (node == null) {
return null
}
if (key < node.key) {
node.left = deleteRecursive(node.left, key)
} else if (key > node.key) {
node.right = deleteRecursive(node.right, key)
} else {
// 노드를 삭제할 경우
if (node.left == null && node.right == null) {
return null
} else if (node.left == null) {
return node.right
} else if (node.right == null) {
return node.left
} else {
// 두 개의 자식이 있는 경우
val successor = findMin(node.right!!)
node.key = successor.key
node.right = deleteRecursive(node.right, successor.key)
}
}
node.updateHeight()
return balance(node)
}
private fun findMin(node: AVLNode): AVLNode {
var current = node
while (current.left != null) {
current = current.left!!
}
return current
}
private fun balance(node: AVLNode?): AVLNode? {
if (node == null) {
return null
}
if (node.balanceFactor() > 1) {
if (node.left?.balanceFactor() ?: 0 < 0) {
node.left = leftRotate(node.left!!)
}
return rightRotate(node)
}
if (node.balanceFactor() < -1) {
if (node.right?.balanceFactor() ?: 0 > 0) {
node.right = rightRotate(node.right!!)
}
return leftRotate(node)
}
return node
}
private fun leftRotate(node: AVLNode): AVLNode {
val newRoot = node.right!!
node.right = newRoot.left
newRoot.left = node
node.updateHeight()
newRoot.updateHeight()
return newRoot
}
private fun rightRotate(node: AVLNode): AVLNode {
val newRoot = node.left!!
node.left = newRoot.right
newRoot.right = node
node.updateHeight()
newRoot.updateHeight()
return newRoot
}
fun printInOrder() {
printInOrderRecursive(root)
println()
}
private fun printInOrderRecursive(node: AVLNode?) {
if (node != null) {
printInOrderRecursive(node.left)
print("${node.key} ")
printInOrderRecursive(node.right)
}
}
fun bfs(node: AVLNode?) {
if (node == null) return
val queue: Queue<Pair<AVLNode, 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 = AVLTree2()
avlTree.insert(3)
avlTree.insert(2)
avlTree.insert(1)
avlTree.insert(7)
avlTree.insert(6)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.insert(5)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.insert(4)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.delete(7)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.delete(5)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.delete(4)
avlTree.bfs(avlTree.root)
println("\n-----------------")
avlTree.delete(6)
avlTree.bfs(avlTree.root)
println("\n-----------------")
}
/*
2(3)
1(1) 6(2)
3(1) 7(1)
-----------------
3(3)
2(2) 6(2)
1(1) 5(1) 7(1)
-----------------
3(4)
2(2) 6(3)
1(1) 5(2) 7(1)
4(1)
-----------------
3(3)
2(2) 5(2)
1(1) 4(1) 6(1)
-----------------
3(3)
2(2) 6(2)
1(1) 4(1)
-----------------
3(3)
2(2) 6(1)
1(1)
-----------------
2(2)
1(1) 3(1)
-----------------
*/
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment