Advertisement
nachtvaohal

tree-split

Feb 13th, 2024
41
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Kotlin 1.57 KB | None | 0 0
  1. fun split(root: Node?, k: Int): List<Node?> {
  2.     // Вычисляем размер левого поддерева с учётом того,
  3.     // что оно может быть пустым.
  4.     val leftSize = if (root?.left == null) 0 else root.left!!.size
  5.     // Если слева ровно k вершин, то искомая вершина - корень консистентного дерева
  6.     if (leftSize == k || k == root?.size) {
  7.         return listOf(root, null)
  8.     }
  9.     // левое поддерево слишком мало
  10.     return if (leftSize < k) {
  11.         // root выше слева
  12.         // поиск в правом поддереве
  13.         val result = split(root!!.right, k - leftSize - 1)
  14.         val rightSubTree = result[0]!!
  15.         val detachedTree = result[1]
  16.         if (detachedTree == null) {
  17.             root.right = rightSubTree.left
  18.             root.recalculateSize()
  19.             rightSubTree.left = null
  20.             rightSubTree.recalculateSize()
  21.             listOf(root, rightSubTree)
  22.         } else {
  23.             root.recalculateSize()
  24.             listOf(root, detachedTree)
  25.         }
  26.     } else {
  27.         // root выше справа
  28.         // поиск в левом поддереве.
  29.         val result = split(root!!.left, k)
  30.         val leftSubTree = result[0]!!
  31.         val rightSubtree = result[1]
  32.         root.left = rightSubtree
  33.         root.recalculateSize()
  34.         listOf(leftSubTree, root)
  35.     }
  36. }
  37.  
  38. fun Node.recalculateSize() {
  39.     size = (left?.size ?: 0) + (right?.size ?: 0) + 1
  40. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement