[자료 구조/알고리즘] 이진 탐색 트리 - 구현

class Node(
    var key: Int,
    var left: Node? = null,
    var right: Node? = null
) {
    fun find(value: Int): Node? = when {
        value < this.key -> left?.find(value)
        value > this.key -> right?.find(value)
        else -> this
    }

    fun insert(value: Int) {
        when {
            value < this.key ->
                if (this.left == null) this.left = Node(value) else this.left?.insert(value)
            value > this.key ->
                if (this.right == null) this.right = Node(value) else this.right?.insert(value)
        }
    }

    fun delete(value: Int) {
        when {
            value < this.key -> searchForDeletion(value, this.left, this)
            value > this.key -> searchForDeletion(value, this.right, this)
            else -> removeNode(this, null)
        }
    }

    fun searchForDeletion(value: Int, node: Node?, parent: Node?) {
        if (node == null) {
            return
        }
        when {
            value < node.key -> searchForDeletion(value, node.left, node)
            value > node.key -> searchForDeletion(value, node.right, node)
            value == node.key -> removeNode(node, parent)
        }
    }

    fun removeNode(node: Node, parent: Node?) {
        // 왼쪽 노드가 있으면
        node.left?.let { left ->
            // 오른쪽 노드도 있으면
            node.right?.let {
                removeNodeWithTwoChilds(node)
            }
            // 왼쪽 노드만 삭제
            ?: removeNodeWithSingleChild(node, left)
        }
        // 왼쪽 노드가 없으면
        ?: run {
            // 오른쪽 노드만 있으면
            node.right?.let {
                right -> removeNodeWithSingleChild(node, right)
            }
            // 리프 노드이면 바로 삭제
            ?: removeLeafNode(node, parent)
        }
    }

    fun removeLeafNode(node: Node, parent: Node?) {
        // 루트 노드가 아닌 경우 리프 노드 삭제 후 부모 노드의 자식 노드 참조 제거
        parent?.let { parent ->
            if (node == parent.left) {
                parent.left = null
            } else if (node == parent.right) {
                parent.right = null
            }
        }
        // 리프 노드이면서 루트 노드인 경우
        ?: throw IllegalStateException("자식 노드가 없는 루트 노드를 삭제할 수 없습니다.")
    }

    fun removeNodeWithTwoChilds(node: Node) {
        // 왼쪽 자식 노드의 하위 노드에서 최대값 탐색
        val left = node.left
        if (left == null) {
            throw IllegalStateException("자식 노드가 두 개인 노드의 왼쪽 자식 노드는 널일 수 없습니다.")
        }

        left.right?.let {
            val maxParent = findMaxChild(left)
            maxParent.right?.let {
                node.key = it.key
                maxParent.right = null
            } ?: throw IllegalStateException("최대 자식 노드를 갖는 노드는 오른쪽 노드를 가져야 합니다.")
        } ?: run {
            node.key = left.key
            node.left = left.left
        }
    }

    fun removeNodeWithSingleChild(parent: Node, child: Node) {
        parent.key = child.key
        parent.left = child.left
        parent.right = child.right
    }

    fun findMaxChild(node: Node): Node {
        // 오른쪽 노드만 탐색
        return node.right?.let {
            r -> r.right?.let {
                findMaxChild(r)
            }
            // 리프 노드면 최대값
            ?: node
        } ?: throw IllegalArgumentException("오른쪽 자식 노드는 널일 수 없습니다.")
    }
}

fun main() {
    val tree = Node(4)
    val keys = arrayOf(8, 15, 21, 3, 7, 2, 5, 10, 2, 3, 4, 6, 11)
    for (key in keys) {
        tree.insert(key)
    }

    println(tree.find(15)?.key ?: 0)
    println(tree.find(15)?.left?.key ?: 0)
    println(tree.find(15)?.right?.key ?: 0)

    tree.delete(15)
    println(tree.find(15)?.key ?: 0)
    println(tree.find(11)?.key ?: 0)
    println(tree.find(11)?.left?.key ?: 0)
    println(tree.find(11)?.right?.key ?: 0)

    tree.insert(13)
    println(tree.find(21)?.key ?: 0)
    println(tree.find(21)?.left?.key ?: 0)

    tree.delete(21)
    println(tree.find(21)?.key ?: 0)
    println(tree.find(11)?.left?.key ?: 0)
    println(tree.find(11)?.right?.key ?: 0)
}