Implement AVL tree

trees/build.gradle.kts

@@ -13,7 +13,7 @@ dependencies {
     testImplementation("org.jetbrains.kotlin:kotlin-test-junit5")
 
     // Use the JUnit 5 integration.
-    testImplementation("org.junit.jupiter:junit-jupiter-engine:5.9.1")
+    testImplementation("org.junit.jupiter:junit-jupiter-engine:5.9.2")
 }
 
 tasks.named<Test>("test") {

trees/src/main/kotlin/bst/AVLTree.kt

@@ -1,3 +1,97 @@
+//GPL-3.0-or-later
+// <Here is a data structure that implements the binary search tree.>
+//This file is part of Trees-3.
+//
+//Trees-3 is free software: you can redistribute and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation; either version 3 of the license or (at your option) any later version.
+//
+//Trees-3 is distributed in the hope that it will be useful, but WITHOUT ANY GUARANTEES; even without an implicit guarantee of merchantability or FITNESS FOR A PARTICULAR PURPOSE. For more information, see the GNU General Public License.
+//
+//You should have obtained a copy of the GNU General Public License with this program. If it is not, see <https://www.gnu.org/licenses/>.
+//    Copyright (C) <2023>  <Nemakin Nikita Antonovich>
+
 package bst
 
-class AVLTree
+import bst.nodes.AVLNode
+import kotlin.math.max
+
+class AVLTree<K: Comparable<K>, V> : BalancingTree<K, V, AVLNode<K, V>>() {
+    override fun initNode(key: K, value: V): AVLNode<K, V> = AVLNode(key, value)
+
+    override fun insertNode(node: AVLNode<K, V>?, key: K, value: V): AVLNode<K, V> {
+        if (node == null) return initNode(key, value)
+        if (key < node.key) {
+            node.left = insertNode(node.left, key, value)
+        } else if (key > node.key) {
+            node.right = insertNode(node.right, key, value)
+        } else {
+            node.value = value
+        }
+        updateHeight(node)
+        return balance(node)
+    }
+
+    private fun getHeight(node: AVLNode<K, V>?): Int {
+        return node?.height ?: -1
+    }
+    private fun updateHeight(node: AVLNode<K, V>) {
+        node.height = max(getHeight(node.left), getHeight(node.right)) + 1
+    }
+    private fun getBalanceFactor(node: AVLNode<K, V>?): Int = when (node) {
+        null -> 0
+        else -> getHeight(node.right) - getHeight(node.left)
+    }
+
+    private fun balance(node: AVLNode<K, V>): AVLNode<K, V> {
+        return when (getBalanceFactor(node)) {
+            -2 -> {
+                if (getBalanceFactor(node.left) == 1) {
+                    node.left = rotateLeft(node.left!!)
+                }
+                return rotateRight(node)
+            }
+            2 -> {
+                if (getBalanceFactor(node.right) == -1) {
+                    node.right = rotateRight(node.right!!)
+                }
+                return rotateLeft(node)
+            }
+            else -> node
+        }
+    }
+
+    override fun removeNode(node: AVLNode<K, V>?, key: K): AVLNode<K, V>? {
+        if (node == null) return null
+        if (key < node.key) {
+            node.left = removeNode(node.left, key)
+        } else if (key > node.key) {
+            node.right = removeNode(node.right, key)
+        } else {
+            if (node.left == null) {
+                return node.right
+            } else if (node.right == null) {
+                return node.left
+            } else {
+                val tmp: AVLNode<K, V> = findMax(node.left)!!
+                node.key = tmp.key
+                node.value = tmp.value
+                node.left = removeNode(node.left, tmp.key)
+            }
+        }
+
+        updateHeight(node)
+        return balance(node)
+    }
+
+    override fun rotateRight(node: AVLNode<K, V>): AVLNode<K, V> {
+        val tmp = super.rotateRight(node)
+        updateHeight(node)
+        updateHeight(tmp)
+        return tmp
+    }
+    override fun rotateLeft(node: AVLNode<K, V>): AVLNode<K, V> {
+        val tmp = super.rotateLeft(node)
+        updateHeight(node)
+        updateHeight(tmp)
+        return tmp
+    }
+}

trees/src/main/kotlin/bst/AbstractBST.kt

@@ -0,0 +1,84 @@
+package bst
+
+import bst.nodes.BinaryNode
+
+abstract class AbstractBST<K: Comparable<K>, V, Self: BinaryNode<K, V, Self>> : Tree<K, V> {
+    internal var rootNode: Self? = null
+    // factory method
+    protected abstract fun initNode(key: K, value: V): Self
+
+    override fun insert(key: K, value: V) {
+        rootNode = insertNode(rootNode, key, value)
+    }
+    protected open fun insertNode(node: Self?, key: K, value: V): Self {
+        if (node == null) return initNode(key, value)
+        if (key < node.key) {
+            node.left = insertNode(node.left, key, value)
+        } else if (key > node.key) {
+            node.right = insertNode(node.right, key, value)
+        } else {
+            node.value = value
+        }
+        return node
+    }
+
+    override fun remove(key: K) {
+        rootNode = removeNode(rootNode, key)
+    }
+    protected open fun removeNode(node: Self?, key: K): Self? {
+        if (node == null) return null
+        if (key < node.key) {
+            node.left = removeNode(node.left, key)
+        } else if (key > node.key) {
+            node.right = removeNode(node.right, key)
+        } else {
+            if (node.left == null) {
+                return node.right
+            } else if (node.right == null) {
+                return node.left
+            } else {
+                val tmp: Self = findMax(node.left)!!
+                node.key = tmp.key
+                node.value = tmp.value
+                node.left = removeNode(node.left, tmp.key)
+            }
+        }
+        return node
+    }
+
+    override fun find(key: K): V? = findNode(rootNode, key)
+    private fun findNode(node: Self?, key: K): V? {
+        return if (node == null) {
+            null
+        } else if (key == node.key) {
+            node.value
+        } else {
+            if (key < node.key)
+                findNode(node.left, key)
+            else
+                findNode(node.right, key)
+        }
+    }
+
+    protected fun findMax(node: Self?): Self? = when {
+        node == null -> null
+        node.right == null -> node
+        else -> findMax(node.right)
+    }
+
+    // for debug purposes only
+    fun printTree() = println(diagram(rootNode))
+    private fun diagram(node: Self?,
+                        top: String = "",
+                        root: String = "",
+                        bottom: String = ""): String {
+        return node?.let {
+            if (node.left == null && node.right == null) {
+                "$root${node.value}\n"
+            } else {
+                diagram(node.right, "$top ", "$top┌──", "$top│ ") +
+                        root + "${node.value}\n" + diagram(node.left, "$bottom│ ", "$bottom└──", "$bottom ")
+            }
+        } ?: "${root}null\n"
+    }
+}

trees/src/main/kotlin/bst/BSTree.kt

@@ -0,0 +1,7 @@
+package bst
+
+import bst.nodes.BSTNode
+
+class BSTree<K: Comparable<K>, V> : AbstractBST<K, V, BSTNode<K, V>>() {
+    override fun initNode(key: K, value: V): BSTNode<K, V> = BSTNode(key, value)
+}

trees/src/main/kotlin/bst/BalancingTree.kt

@@ -1,15 +1,15 @@
 package bst
 
-import bst.nodes.TreeNode
+import bst.nodes.BinaryNode
 
-abstract class BalancingTree<K : Comparable<K>, V> : BinarySearchTree<K, V>() {
-    protected fun <SpecNode : TreeNode<K, V, SpecNode>> rotateLeft(node: SpecNode): SpecNode {
+abstract class BalancingTree<K: Comparable<K>, V, Self: BinaryNode<K, V, Self>>: AbstractBST<K, V, Self>() {
+    protected open fun rotateLeft(node: Self): Self {
         val right = node.right
         node.right = right?.left
         right?.left = node
         return right!!
     }
-    protected fun <SpecNode : TreeNode<K, V, SpecNode>> rotateRight(node: SpecNode): SpecNode {
+    protected open fun rotateRight(node: Self): Self {
         val left = node.left
         node.left = left?.right
         left?.right = node

trees/src/main/kotlin/bst/RedBlackTree.kt

@@ -2,9 +2,7 @@ package bst
 
 import bst.nodes.RBTNode
 
-class RedBlackTree<K : Comparable<K>, V> : BalancingTree<K, V>() {
-    private var root: RBTNode<K, V>? = null
-
+class RedBlackTree<K : Comparable<K>, V> : BalancingTree<K, V, RBTNode<K, V>>() {
     private fun isRed(node: RBTNode<K, V>?): Boolean {
         return node?.red == true
     }
@@ -14,12 +12,12 @@ class RedBlackTree<K : Comparable<K>, V> : BalancingTree<K, V>() {
     }
 
     private fun insertNode(key: K, value: V): RBTNode<K, V> {
-        if (root == null) {
-            root = RBTNode(key, value, true)
-            return root!!
+        if (rootNode == null) {
+            rootNode = RBTNode(key, value, true)
+            return rootNode!!
         }
 
-        var node = root
+        var node = rootNode
         var parent: RBTNode<K, V>? = null
 
         while (node != null) {
@@ -81,7 +79,7 @@ class RedBlackTree<K : Comparable<K>, V> : BalancingTree<K, V>() {
             }
         }
 
-        root!!.red = true
+        rootNode!!.red = true
         return current
     }
 
@@ -92,4 +90,8 @@ class RedBlackTree<K : Comparable<K>, V> : BalancingTree<K, V>() {
     private fun removeNode(node: RBTNode<K, V>): RBTNode<K, V> {
         TODO()
     }
+
+    override fun initNode(key: K, value: V): RBTNode<K, V> {
+        TODO("Not yet implemented")
+    }
 }

trees/src/main/kotlin/bst/Tree.kt

@@ -3,6 +3,5 @@ package bst
 interface Tree<K : Comparable<K>, V> {
     fun insert(key: K, value: V)
     fun remove(key: K)
-    fun find(key: K): Boolean
-    fun clear()
+    fun find(key: K): V?
 }

trees/src/main/kotlin/bst/nodes/AVLNode.kt

@@ -1,3 +1,8 @@
 package bst.nodes
 
-class AVLNode
+class AVLNode<K: Comparable<K>, V>(
+    key: K,
+    value: V
+) : BinaryNode<K, V, AVLNode<K, V>>(key, value) {
+    var height: Int = 0
+}