Kotlin Red Black Tree. For a programming challenge. Instead of a key-value store, it returns the number of items added, whose value is greater than than the given key.

GreaterThanBST.kt

class GreaterThanBST {

    private var root : Node? = null

    fun add(key : Int) {
        root = put(root, key)
    }

    operator fun get(key : Int) : Int = get(root, key)

    companion object {

        private fun get(node: Node?, key: Int) : Int {
            var value = 0
            var n = node
            while (n != null) {
                when {
                    key < n.key -> {
                        value += (n.gtNode?.value ?: 0) + n.instances
                        n = n.ltNode
                    }
                    key > n.key -> {
                        n = n.gtNode
                    }
                    else -> return (n.gtNode?.value ?: 0) + value
                }
            }
            return 0
        }

        private fun put(node: Node?, key: Int) : Node {
            return node?.let {
                when {
                    key < node.key ->
                        node.copy(ltNode = put(node.ltNode, key), value = node.value + 1)
                    key > node.key ->
                        node.copy(gtNode = put(node.gtNode, key), value = node.value + 1)
                    else -> node.addInstance()
                }.run {
                    if (rightBlackLeftRed) rotateLeft()
                    else this
                }.run {
                    if (consecutiveLeftRed) rotateRight()
                    else this
                }.run {
                    if (bothChildrenRed) flipColors()
                    else this
                }
            } ?: Node(key)
        }
    }

    private enum class Color {
        Red, Black;
        fun other() = if (this == Red) Black else Red
    }

    private data class Node(val key : Int,
                            val value: Int = 1,
                            val instances: Int = 1,
                            val size: Int = 1,
                            val gtNode: Node? = null,
                            val ltNode: Node? = null,
                            val color: Color = Color.Red) {

        fun rotateRight() = ltNode!!.copy(
                gtNode = this.becomeGtNodeInRightRotation(ltNode.gtNode),
                color = this.color,
                size = this.size,
                value = this.value
        )

        private fun becomeGtNodeInRightRotation( newLtNode : Node? ) = copy(
                color = Color.Red,
                ltNode = newLtNode,
                size = 1 + sizeOf(newLtNode) + sizeOf(gtNode),
                value = instances + valueOf(newLtNode) + valueOf(gtNode)
        )

        fun rotateLeft() = gtNode!!.copy(
                ltNode = this.becomeLtNodeInLeftRotation(newGtNode = gtNode.ltNode),
                color = this.color,
                size = this.size,
                value = this.value
        )

        private fun becomeLtNodeInLeftRotation(newGtNode : Node?) = copy(
                color = Color.Red,
                gtNode = newGtNode,
                value = instances + valueOf(newGtNode) + valueOf(ltNode),
                size = 1 + sizeOf(newGtNode) + sizeOf(ltNode)
        )

        fun flipColors() = copy(color = color.other(),
                ltNode = ltNode!!.copy(color = ltNode.color.other()),
                gtNode = gtNode!!.copy(color = gtNode.color.other()))

        val rightBlackLeftRed get() = red(gtNode) && black(ltNode)
        val consecutiveLeftRed get() = red(ltNode) && red(ltNode?.ltNode)
        val bothChildrenRed get() = red(ltNode) && red(gtNode)
        fun addInstance() = copy(instances = instances + 1, value = value + 1)
        private fun sizeOf(node : Node?) = node?.size ?: 0

        companion object {
            fun valueOf(node : Node?) = node?.value ?: 0
            fun instancesOf(node : Node?) = node?.instances ?: 0
            fun red(node : Node?) = (node?.color ?: Color.Black) == Color.Red
            fun black(node : Node?) = (node?.color ?: Color.Black) == Color.Black
        }
    }
}