Implement Red Black Tree

Author: Romanow

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

@@ -18,11 +18,69 @@ class RedBlackTree {
         return null
     }
 
+    fun put(key: Int, value: Int = key) {
+        root = insert(root, key, value)
+        root?.color = false
+    }
+
+    fun remove(key: Int) {
+
+    }
+
+    private fun insert(n: TreeNode?, key: Int, value: Int): TreeNode {
+        var node = n ?: 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
+        }
+
+        if (isRed(node.right) && !isRed(node.left)) {
+            node = rotateLeft(node)
+        }
+        if (isRed(node.left) && isRed(node.left?.left)) {
+            node = rotateRight(node)
+        }
+        if (isRed(node.left) && isRed(node.right)) {
+            flipColors(node)
+        }
+
+        return node
+    }
+
+    private fun isRed(node: TreeNode?): Boolean {
+        return node?.color == true
+    }
+
+    private fun rotateLeft(node: TreeNode): TreeNode {
+        val right = node.right!!
+        node.right = right.left
+        right.left = node
+        right.color = node.color
+        node.color = true
+        return right
+    }
+
+    private fun rotateRight(node: TreeNode): TreeNode {
+        val left = node.left!!
+        node.left = left.right
+        left.right = node
+        left.color = node.color
+        node.color = true
+        return left
+    }
+
+    private fun flipColors(node: TreeNode) {
+        node.color = true
+        node.left?.color = false
+        node.right?.color = false
+    }
+
     internal class TreeNode(
         override var key: Int,
         override var value: Int,
         var color: Boolean = false,
-        override var left: TreeNode?,
-        override var right: TreeNode?
+        override var left: TreeNode? = null,
+        override var right: TreeNode? = null
     ) : TNode<TreeNode>
 }

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

@@ -12,7 +12,7 @@ class BinarySearchListTest {
 
     @ArgumentsSource(ValueProvider::class)
     @ParameterizedTest(name = "#{index} – Search {1} in tree {0}")
-    fun flatten(list: List<Int>, target: Int, result: Int?) {
+    fun test(list: List<Int>, target: Int, result: Int?) {
         val obj = BinarySearchList()
         assertThat(obj.find(list, target)).isEqualTo(result)
     }

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

@@ -12,7 +12,7 @@ class BinarySearchTreeTest {
 
     @ArgumentsSource(ValueProvider::class)
     @ParameterizedTest(name = "#{index} – Search {1} in tree {0}")
-    fun flatten(items: List<Int>, key: Int, result: Int?) {
+    fun test(items: List<Int>, key: Int, result: Int?) {
         val obj = BinarySearchTree()
         items.forEach { obj.put(it) }
         assertThat(obj.get(key)).isEqualTo(result)

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

@@ -13,7 +13,7 @@ class BreadthFirstSearchTest {
 
     @ArgumentsSource(ValueProvider::class)
     @ParameterizedTest(name = "#{index} – Bfs {0} is {1}")
-    fun flatten(items: List<Int?>, result: List<Int>) {
+    fun test(items: List<Int?>, result: List<Int>) {
         val root = buildTreeFromList(items)
         // printTree(root)
 

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

@@ -13,7 +13,7 @@ class InOrderDeepFirstSearchTest {
 
     @ArgumentsSource(ValueProvider::class)
     @ParameterizedTest(name = "#{index} – In-order dfs {0} is {1}")
-    fun flatten(items: List<Int?>, result: List<Int>) {
+    fun test(items: List<Int?>, result: List<Int>) {
         val root = buildTreeFromList(items)
         // printTree(root)
 

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

@@ -13,7 +13,7 @@ class PostOrderDeepFirstSearchTest {
 
     @ArgumentsSource(ValueProvider::class)
     @ParameterizedTest(name = "#{index} – Post-order dfs {0} is {1}")
-    fun flatten(items: List<Int?>, result: List<Int>) {
+    fun test(items: List<Int?>, result: List<Int>) {
         val root = buildTreeFromList(items)
         // printTree(root)
 

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

@@ -13,7 +13,7 @@ class PreOrderDeepFirstSearchTest {
 
     @ArgumentsSource(ValueProvider::class)
     @ParameterizedTest(name = "#{index} – Pre-order dfs {0} is {1}")
-    fun flatten(items: List<Int?>, result: List<Int>) {
+    fun test(items: List<Int?>, result: List<Int>) {
         val root = buildTreeFromList(items)
         // printTree(root)
 

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

@@ -0,0 +1,30 @@
+package ru.romanow.algorithms
+
+import org.assertj.core.api.Assertions.assertThat
+import org.junit.jupiter.api.extension.ExtensionContext
+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 java.util.stream.Stream
+
+class RedBlackTreeTest {
+
+    @ArgumentsSource(ValueProvider::class)
+    @ParameterizedTest(name = "#{index} – Search {1} in tree {0}")
+    fun test(items: List<Int>, key: Int, result: Int?) {
+        val obj = RedBlackTree()
+        items.forEach { obj.put(it) }
+        assertThat(obj.get(key)).isEqualTo(result)
+        obj.remove(key)
+        assertThat(obj.get(key)).isNull()
+    }
+
+    internal class ValueProvider : ArgumentsProvider {
+        override fun provideArguments(context: ExtensionContext): Stream<Arguments> =
+            Stream.of(
+                Arguments.of(listOf(10, 5, 15, 3, 7, 12, 20), 20, 20),
+                Arguments.of(listOf(10, 5, 15, 3, 7, 12, 20), 0, null)
+            )
+    }
+}