remove-node

import java.util.*

// <template>
data class Node(var left: Node?, var right: Node?, var value: Int)
// <template>


/**
 *
 *
 * ПРИНЦИП РАБОТЫ
 *
 *
 * ДОКАЗАТЕЛЬСТВО КОРРЕКТНОСТИ
 * На место удаляемой вершины можно поставить
 * * самую правую вершину в левом поддереве (наибольшую из поддерева)
 * * самую левую вершину в правом поддереве
 *
 * ВРЕМЕННАЯ СЛОЖНОСТЬ
 *
 * ПРОСТРАНСТВЕННАЯ СЛОЖНОСТЬ
 *
 * @param root корень дерева
 * @param key значение, которое нужно удалить
 *
 * @return корень изменённого дерева
 */
fun remove(root: Node?, key: Int): Node? {
    if (root == null) {
        return null
    }
    val parents = Stack<Node>()
    val toRemove = findNodeByKey(root, key, parents)
    // 0. Удаляемый узел не найден
    if (toRemove == null) {
        return root
    }
    // 1. Если дерево состояло из одной вершины, то после её удаления дерева не останется
    if (toRemove.value == root.value && root.isLeaf()) {
        return null
    }
    // 2. Если мы удаляем лист, то дерево останется одним деревом и не распадётся на части
    val parentOfNodeToRemove = if (parents.size == 0) null else parents.pop()
    return if (parentOfNodeToRemove != null) {
        val replacingNode = findReplacementNode(toRemove)
        if (replacingNode == toRemove) {
            if (parentOfNodeToRemove.value > replacingNode.value) {
                parentOfNodeToRemove.left = null
            } else {
                parentOfNodeToRemove.right = null
            }
        } else {
            replacingNode.right = toRemove.right
            if (parentOfNodeToRemove.value > replacingNode.value) {
                parentOfNodeToRemove.left = replacingNode
            } else {
                parentOfNodeToRemove.right = replacingNode
            }
        }
        root
    } else {
        if (toRemove.left != null) {
            val maxInLeftTree = findReplacementNode(toRemove)
            maxInLeftTree.right = toRemove.right
            maxInLeftTree.left = toRemove.left
            maxInLeftTree
        } else {
            toRemove.right
        }
    }
}

/**
 * Найти узел по ключу с сохранением родителя найденной ноды в переданном стеке.
 *
 * @param root корень дерева
 * @param key ключ, по которому осуществляется поиск
 * @param parents стек, в котором сохраняется родитель найденного узла
 *
 * @return найденный узел
 * @return null, если узел не был найден
 */
fun findNodeByKey(root: Node?, key: Int, parents: Stack<Node>): Node? {
    if (root == null) {
        parents.clear()
        return null
    }
    if (root.value == key) {
        return root
    }
    parents.push(root)
    return if (key < root.value) {
        findNodeByKey(root.left, key, parents)
    } else {
        findNodeByKey(root.right, key, parents)
    }
}

/**
 * Найти наибольший элемент в левом поддереве удаляемого узла.
 *
 * @param toDelete удаляемый элемент
 * @return наибольший (самый правый) элемент в левом поддереве удаляемого элемента [toDelete], либо сам [toDelete]
 */
fun findReplacementNode(toDelete: Node): Node {
    if (toDelete.isLeaf()) {
        return toDelete
    }

    val maxRight = if (toDelete.left == null) {
        extractMax(toDelete.right!!)
    } else {
        extractMax(toDelete.left!!)
    }
    if (toDelete.right == maxRight) {
        toDelete.right = null
    }
    return maxRight
}

/**
 * Извлечь наибольший узел из переданного узла.
 *
 * @param node поддерево, в котором нужно найти наибольший узел
 * @return наибольший узел или само поддерево, если оно является единственным узлом
 */
fun extractMax(node: Node): Node {
    val max = if (node.right != null) {
        extractMax(node.right!!)
    } else {
        node
    }
    // удаление наибольшего узла из его предка
    if (node.right == max) {
        node.right = null
    }
    return max
}

private fun Node.isLeaf(): Boolean {
    return right == null && left == null
}