AVL-tree, unit-tested and random tested with op sequences of length up to 2e6

AVL.kt

package algo

import toInt
import kotlin.math.absoluteValue


class AVL<K : Comparable<K>, V : Any> {
    private var Node?.diff
        get() = this?._diff ?: 0
        set(value) {
            if (this != null) _diff = value
        }

    private inner class Node(val key: K, var value: V) {
        private val child = arrayOfNulls<Node>(2)
        //TODO private var parent: Node? = null //needed for effective (non-recursive) traverse

        @Suppress("PropertyName")
        internal var _diff = 0

        private fun which(queryKey: K) = (key < queryKey).toInt()



        fun get(key_: K): Node? = if (key == key_) this else child[which(key_)]?.get(key_)



        private val relLeftIndex get() = (diff < 0).toInt()
        private var relLeft
            get() = child[relLeftIndex]
            set(v) {
                child[relLeftIndex] = v
            }
        private var relRight
            get() = child[1 - relLeftIndex];
            set(v) {
                child[1 - relLeftIndex] = v
            }

        private fun smallRotate(): Node {
            val b = relRight!! //rotation is left according to relative sides
            val middle = b.relLeft
            b.relLeft = this
            this.relRight = middle

            diff = b.diff

            return b
        }

        private fun bigRotate(): Node {
            relRight = relRight!!.smallRotate() //rotation is left according to relative sides
            return smallRotate()
        }

        /**
         * @param value_ null to delete key
         * @returns *new root of subtree
         * @throws NullPointerException if there is no such key to delete
         * */
        fun baseSet(key_: K, value_: V?): Node? {
            if (key == key_) {
                return if (value_ == null) {
                    return child[0] ?: child[1]
                } else {
                    value = value_
                    this
                }
            }

            child[which(key_)] = child[which(key_)]?.baseSet(key_, value_) ?: Node(
                key_,
                value_!! //this place will throw NPE if there is no such element to delete
            )

            if (diff.absoluteValue == 2) {
                if (relRight.diff == -diff / 2) //diff of relatively right node is 1
                    relRight = relRight!!.smallRotate()
                return smallRotate()
            }

            return this
        }



        fun toList(list: MutableList<V>): List<V> = list.apply {
            child[0]?.flatten(list)
            add(value)
            child[1]?.flatten(list)
        }
    }

    private var root: Node? = null
    var size = 0
        private set

    operator fun set(key: K, value: V) {
        ++size
        root = if (root == null)
            Node(key, value)
        else
            root!!.baseSet(key, value)
    }

    /** @throws NullPointerException if there is no such element */
    fun delete(key: K) {
        --size
        root = root!!.baseSet(key, null)
    }

    /** @return null if there is no such key */
    operator fun get(key: K) = root?.get(key)?.value

    fun toList() = root?.toList(mutableListOf()) ?: listOf()
}