add red & black BST

src/main/kotlin/it/norangeb/algorithms/datastructures/dictionary/RedBlackBST.kt

@@ -25,66 +25,182 @@
 
 package it.norangeb.algorithms.datastructures.dictionary
 
+import arrow.core.None
 import arrow.core.Option
+import arrow.core.toOption
 
 class RedBlackBST<K : Comparable<K>, V> : OrderedDictionary<K, V> {
-    override fun get(key: K): Option<V> {
+    private var root: Option<Node<K, V>> = None
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+
+    override fun get(key: K): Option<V> = get(root, key).map { it.value }
+
+    private fun get(
+        node: Option<Node<K, V>>,
+        key: K
+    ): Option<Node<K, V>> = node.flatMap {
+        when {
+            key == it.key -> node
+            key > it.key -> get(it.right, key)
+            else -> get(it.left, key)
+        }
     }
 
     override fun set(key: K, value: V) {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        root = set(root, key, value)
     }
 
+    private fun set(
+        node: Option<Node<K, V>>,
+        key: K,
+        value: V
+    ): Option<Node<K, V>> = normalize(
+        node.fold(
+            { Node(key, value) },
+            {
+                when {
+                    key == it.key -> it.clone(value = value)
+                    key > it.key -> it.clone(right = set(it.right, key, value))
+                    else -> it.clone(left = set(it.left, key, value))
+                }
+            }
+        )
+    ).toOption()
+
+    private fun normalize(node: Node<K, V>): Node<K, V> {
+        var resultNode = node
+
+        if (isRightRed(resultNode))
+            resultNode = rotateLeft(resultNode)
+
+        if (isTwoRedInRow(resultNode))
+            resultNode = rotateRight(resultNode)
+
+        if (isBothRed(resultNode))
+            resultNode = colorFlip(resultNode)
+
+        return resultNode
+    }
+
+    private fun rotateRight(node: Node<K, V>): Node<K, V> = node
+        .left.fold(
+        { node },
+        {
+            it.clone(
+                color = node.color,
+                right = node.clone(color = Color.RED, left = it.right).toOption()
+            )
+        }
+    )
+
+    private fun rotateLeft(node: Node<K, V>): Node<K, V> = node
+        .right.fold(
+        { node },
+        {
+            it.clone(
+                color = node.color,
+                left = node.clone(color = Color.RED, right = it.left).toOption()
+            )
+        }
+    )
+
+    private fun colorFlip(node: Node<K, V>): Node<K, V> = node.clone(
+        color = Color.RED,
+        left = node.left.map { it.clone(color = Color.BLACK) },
+        right = node.right.map { it.clone(color = Color.BLACK) }
+    )
+
+    private fun isRightRed(node: Node<K, V>): Boolean = isRed(node.right) &&
+            !isRed(node.left)
+
+    private fun isTwoRedInRow(node: Node<K, V>): Boolean = isRed(node.left) &&
+            isRed(node.left.flatMap { it.left })
+
+    private fun isBothRed(node: Node<K, V>): Boolean = isRed(node.left) &&
+            isRed(node.right)
+
+    private fun isRed(node: Option<Node<K, V>>): Boolean = node
+        .fold(
+            { false },
+            { it.isRed() }
+        )
+
     override fun delete(key: K) {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun contains(key: K): Boolean {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun isEmpty(): Boolean {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun size(): Int {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun max(): Option<K> {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun min(): Option<K> {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun floor(key: K): Option<K> {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun ceiling(key: K): Option<K> {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun select(pos: Int): Option<K> {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun rank(key: K): Int {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
-    override fun <R> preOrder(transform: (K) -> R) {
+    override fun <R> preOrder(transform: (K) -> R) = preOrder(root, transform)
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+
+    private fun <R> preOrder(node: Option<Node<K, V>>, transform: (K) -> R) {
+        node.map {
+            transform(it.key)
+            preOrder(it.left, transform)
+            preOrder(it.right, transform)
+        }
     }
 
     override fun <R> inOrder(transform: (K) -> R) {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
     }
 
     override fun <R> postOrder(transform: (K) -> R) {
-        TODO("not implemented") // To change body of created functions use File | Settings | File Templates.
+        TODO("the same as bst")
+    }
+
+    data class Node<K, V>(
+        val key: K,
+        val value: V,
+        val color: Color = Color.RED,
+        val left: Option<Node<K, V>> = None,
+        val right: Option<Node<K, V>> = None
+    ) {
+        fun clone(
+            key: K = this.key,
+            value: V = this.value,
+            color: Color = this.color,
+            left: Option<Node<K, V>> = this.left,
+            right: Option<Node<K, V>> = this.right
+        ): Node<K, V> = Node(key, value, color, left, right)
+
+        fun isRed() = color == Color.RED
+    }
+
+    enum class Color {
+        RED, BLACK
     }
 }

src/test/kotlin/it/norangeb/algorithms/datastructures/dictionary/RedBlackBSTTest.kt

@@ -0,0 +1,164 @@
+/*
+ * MIT License
+ *
+ * Copyright (c) 2019 norangebit
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ */
+
+package it.norangeb.algorithms.datastructures.dictionary
+
+import arrow.core.None
+import arrow.core.Some
+import org.amshove.kluent.`should equal`
+import org.spekframework.spek2.Spek
+import org.spekframework.spek2.style.gherkin.Feature
+
+object RedBlackBSTTest : Spek({
+    Feature("red black tree") {
+        Scenario("get and set") {
+            val tree by memoized { RedBlackBST<Int, Int>() }
+
+            When("add 1, 3, 4, 6, 8 into tree") {
+                tree[1] = 1
+                tree[3] = 3
+                tree[4] = 4
+                tree[6] = 6
+                tree[8] = 9
+            }
+
+            Then("tree of 1 should equal 1") {
+                tree[1] `should equal` Some(1)
+            }
+
+            Then("tree of 2 should equal None") {
+                tree[2] `should equal` None
+            }
+
+            Then("tree of 6 should equal 6") {
+                tree[6] `should equal` Some(6)
+            }
+
+            Then("tree of 8 should equal 9") {
+                tree[8] `should equal` Some(9)
+            }
+
+            When("edit 8") {
+                tree[8] = 8
+            }
+
+            Then("tree of 8 should equal 8") {
+                tree[8] `should equal` Some(8)
+            }
+        }
+
+        Scenario("ordered balance") {
+            val tree by memoized { RedBlackBST<Int, Int>() }
+
+            Given("a 1, 2, 3, 4, 5, 6, 7 tree") {
+                tree[1] = 1
+                tree[2] = 2
+                tree[3] = 3
+                tree[4] = 4
+                tree[5] = 5
+                tree[6] = 6
+                tree[7] = 7
+            }
+
+            val result = mutableListOf<Int>()
+
+            When("compute the preOrder traverse") {
+                tree.preOrder { result.add(it) }
+            }
+
+            Then("result should equal listOf(4, 2, 1, 3, 6, 5, 7") {
+                result `should equal` listOf(4, 2, 1, 3, 6, 5, 7)
+            }
+        }
+
+        Scenario("reverse balance") {
+            val tree by memoized { RedBlackBST<Int, Int>() }
+
+            Given("a 7, 6, 5, 4, 3, 2, 1 tree") {
+                tree[7] = 7
+                tree[6] = 6
+                tree[5] = 5
+                tree[4] = 4
+                tree[3] = 3
+                tree[2] = 2
+                tree[1] = 1
+            }
+
+            val result = mutableListOf<Int>()
+
+            When("compute the preOrder traverse") {
+                tree.preOrder { result.add(it) }
+            }
+
+            Then("result should equal listOf(4, 2, 1, 3, 6, 5, 7") {
+                result `should equal` listOf(4, 2, 1, 3, 6, 5, 7)
+            }
+        }
+
+        Scenario("mix balance") {
+            val tree by memoized { RedBlackBST<Int, Int>() }
+
+            Given("a 1, 2, 3, 4, 5, 6, 7 tree") {
+                tree[6] = 6
+                tree[2] = 2
+                tree[4] = 4
+                tree[1] = 1
+                tree[3] = 3
+                tree[7] = 7
+                tree[5] = 5
+            }
+
+            val result = mutableListOf<Int>()
+
+            When("compute the preOrder traverse") {
+                tree.preOrder { result.add(it) }
+            }
+
+            Then("result should equal listOf(4, 2, 1, 3, 6, 5, 7") {
+                result `should equal` listOf(4, 2, 1, 3, 6, 5, 7)
+            }
+
+            When("add 8, 9, 10, 12, 13, 14, 15") {
+                tree[8] = 8
+                tree[9] = 9
+                tree[10] = 10
+                tree[11] = 11
+                tree[12] = 12
+                tree[13] = 13
+                tree[14] = 14
+                tree[15] = 15
+            }
+
+            When("compute the preOrder traverse") {
+                result.clear()
+                tree.preOrder { result.add(it) }
+            }
+
+            Then("result should equal listOf(8, 4, 2, 1, 3, 6, 5, 7, 12, 10, 9, 11, 14, 13, 15)") {
+                result `should equal` listOf(8, 4, 2, 1, 3, 6, 5, 7, 12, 10, 9, 11, 14, 13, 15)
+            }
+        }
+    }
+})