Add main operations to BST

Author: Romanow

src/main/kotlin/ru/romanow/algorithms/BinarySearchTree.kt

@@ -3,12 +3,69 @@ package ru.romanow.algorithms
 import ru.romanow.models.TreeNode
 
 class BinarySearchTree {
-    fun find(root: TreeNode?, target: Int): TreeNode? {
-        return when {
-            root == null -> null
-            target == root.value -> root
-            target < root.value -> find(root.left, target)
-            else -> find(root.right, target)
+    internal var root: TreeNode? = null
+
+    fun get(key: Int): Int? {
+        var current = root
+        while (current != null) {
+            current = when {
+                key < current.key -> current.left
+                key > current.key -> current.right
+                else -> return current.value
+            }
+        }
+        return null
+    }
+
+    fun put(key: Int, value: Int = key) {
+        root = insert(root, key, value)
+    }
+
+    fun remove(key: Int) {
+        root = remove(root, key)
+    }
+
+    private fun insert(node: TreeNode?, key: Int, value: Int = key): TreeNode {
+        if (node == null) {
+            return TreeNode(key, value)
+        }
+        when {
+            key < node.key -> node.left = insert(node.left, key, value)
+            key > node.key -> node.right = insert(node.right, key, value)
+            else -> node.value = value
+        }
+        return node
+    }
+
+    private fun remove(node: TreeNode?, key: Int): TreeNode? {
+        if (node == null) {
+            return null
+        }
+        when {
+            key < node.key -> node.left = remove(node.left, key)
+            key > node.key -> node.right = remove(node.right, key)
+            else -> {
+                if (node.left == null) {
+                    return node.right
+                }
+                if (node.right == null) {
+                    return node.left
+                }
+
+                val minNode = findMin(node.right!!)
+                node.key = minNode.key
+                node.value = minNode.value
+                node.right = remove(node.right, minNode.key)
+            }
+        }
+        return node
+    }
+
+    private fun findMin(node: TreeNode): TreeNode {
+        var current = node
+        while (current.left != null) {
+            current = current.left!!
         }
+        return current
     }
 }

src/main/kotlin/ru/romanow/algorithms/RedBlackTree.kt

@@ -0,0 +1,28 @@
+package ru.romanow.algorithms
+
+import ru.romanow.models.TNode
+
+class RedBlackTree {
+    private var root: TreeNode? = null
+
+    fun get(key: Int): Int? {
+        var current = root
+        while (current != null) {
+            current = when {
+                key < current.key -> current.left
+                key > current.key -> current.right
+                else -> return current.value
+            }
+        }
+        return null
+    }
+
+    internal class TreeNode(
+        override var key: Int,
+        override var value: Int,
+        var color: Boolean = false,
+        override var left: TreeNode?,
+        override var right: TreeNode?
+    ) : TNode<TreeNode>
+
+}

src/main/kotlin/ru/romanow/models/ListNode.kt

@@ -1,18 +1,6 @@
 package ru.romanow.models
 
-class ListNode {
-    var value: Int? = null
-    var next: ListNode? = null
-
-    constructor(value: Int) {
-        this.value = value
-    }
-
-    constructor(value: Int, next: ListNode?) {
-        this.value = value
-        this.next = next
-    }
-
+open class ListNode(var value: Int?, var next: ListNode? = null) {
     override fun toString(): String {
         return "ListNode(value=$value, next=${next?.value})"
     }

src/main/kotlin/ru/romanow/models/TreeNode.kt

@@ -1,25 +1,37 @@
 package ru.romanow.models
 
-import java.util.LinkedList
-import kotlin.collections.ArrayList
+import java.util.*
 
-data class TreeNode(var value: Int) {
-    var left: TreeNode? = null
-    var right: TreeNode? = null
+interface TNode<T : TNode<T>> {
+    var key: Int
+    var value: Int
+    var left: T?
+    var right: T?
 }
 
-fun printTree(node: TreeNode?, level: Int = 0) {
+class TreeNode(
+    override var key: Int,
+    override var value: Int = key,
+    override var left: TreeNode? = null,
+    override var right: TreeNode? = null
+) : TNode<TreeNode> {
+    override fun toString(): String {
+        return "TreeNode(key=$key, value=$value, left=${left?.key}, right=${right?.key})"
+    }
+}
+
+fun printTree(node: TNode<*>?, level: Int = 0) {
     if (node == null) return
     print("  ".repeat(level))
-    println("├─ ${node.value}")
+    println("├─ $node")
     printTree(node.left, level + 1)
     printTree(node.right, level + 1)
 }
 
-fun buildListFromTree(root: TreeNode?): List<Int?> {
+fun buildListFromTree(root: TNode<*>?): List<Int?> {
     val list = ArrayList<Int?>()
     if (root != null) {
-        val queue = LinkedList<TreeNode?>()
+        val queue = LinkedList<TNode<*>?>()
         queue.add(root)
         while (queue.isNotEmpty()) {
             val head = queue.pop()
@@ -38,7 +50,7 @@ fun buildListFromTree(root: TreeNode?): List<Int?> {
 fun buildTreeFromList(values: List<Int?>, index: Int = 0): TreeNode? {
     var node: TreeNode? = null
     if (index < values.size && values[index] != null) {
-        node = TreeNode(value = values[index]!!)
+        node = TreeNode(key = values[index]!!)
         node.left = buildTreeFromList(values, 2 * index + 1)
         node.right = buildTreeFromList(values, 2 * index + 2)
     }

src/test/kotlin/ru/romanow/algorithms/BinarySearchTreeTest.kt

@@ -6,17 +6,18 @@ import org.junit.jupiter.params.ParameterizedTest
 import org.junit.jupiter.params.provider.Arguments
 import org.junit.jupiter.params.provider.ArgumentsProvider
 import org.junit.jupiter.params.provider.ArgumentsSource
-import ru.romanow.models.buildTreeFromList
 import java.util.stream.Stream
 
 class BinarySearchTreeTest {
 
     @ArgumentsSource(ValueProvider::class)
     @ParameterizedTest(name = "#{index} – Search {1} in tree {0}")
-    fun flatten(items: List<Int>, target: Int, result: Int?) {
-        val root = buildTreeFromList(items)
+    fun flatten(items: List<Int>, key: Int, result: Int?) {
         val obj = BinarySearchTree()
-        assertThat(obj.find(root, target)?.value).isEqualTo(result)
+        items.forEach { obj.put(it) }
+        assertThat(obj.get(key)).isEqualTo(result)
+        obj.remove(key)
+        assertThat(obj.get(key)).isNull()
     }
 
     internal class ValueProvider : ArgumentsProvider {