Data update

Task/AVL-tree/Swift/avl-tree.swift

@@ -0,0 +1,243 @@
+import Foundation
+
+// MARK: - AVL Tree ---------------------------------------------------------
+
+final class AVLTree {
+    // ---------- Node ------------------------------------------------------
+    private class Node {
+        var key: Int
+        var balance: Int = 0
+        var height: Int = 0
+
+        var left: Node?
+        var right: Node?
+        weak var parent: Node?          // weak to avoid retain cycles
+
+        init(key: Int, parent: Node?) {
+            self.key = key
+            self.parent = parent
+        }
+    }
+
+    // ---------- Root -------------------------------------------------------
+    private var root: Node?
+
+    // ---------- Public API -------------------------------------------------
+
+    /// Inserts `key`. Returns `true` if the key was added, `false` if it already existed.
+    @discardableResult
+    func insert(_ key: Int) -> Bool {
+        // empty tree → new root
+        guard let rootNode = root else {
+            root = Node(key: key, parent: nil)
+            return true
+        }
+
+        var n: Node? = rootNode
+        while let cur = n {
+            if cur.key == key { return false }          // duplicate
+
+            let goLeft = key < cur.key
+            let parent = cur
+            n = goLeft ? cur.left : cur.right
+
+            // we have found the empty spot → insert
+            if n == nil {
+                let newNode = Node(key: key, parent: parent)
+                if goLeft {
+                    parent.left = newNode
+                } else {
+                    parent.right = newNode
+                }
+                rebalance(parent)                       // fix AVL balance upwards
+                break
+            }
+        }
+        return true
+    }
+
+    /// Deletes `key` if it exists. Does nothing when the key is not present.
+    func delete(_ key: Int) {
+        var current = root
+        while let node = current {
+            if key == node.key {
+                delete(node)                // internal helper that really removes the node
+                return
+            }
+            current = (key < node.key) ? node.left : node.right
+        }
+        // key not found → nothing to do
+    }
+
+    /// Prints the balance factor of every node in‑order.
+    func printBalance() {
+        printBalance(node: root)
+    }
+
+    // ---------- Private helpers -------------------------------------------
+
+    /// Removes `node` from the tree (used by the public `delete(_:)` above).
+    private func delete(_ node: Node) {
+        // ----- 1️⃣ leaf node ------------------------------------------------
+        if node.left == nil && node.right == nil {
+            if let parent = node.parent {
+                if parent.left === node { parent.left = nil }
+                else { parent.right = nil }
+                rebalance(parent)
+            } else {
+                root = nil                     // tree becomes empty
+            }
+            return
+        }
+
+        // ----- 2️⃣ node has a left subtree → replace with predecessor -----
+        if let left = node.left {
+            var predecessor = left
+            while let r = predecessor.right { predecessor = r }
+            node.key = predecessor.key
+            delete(predecessor)
+        }
+        // ----- 3️⃣ otherwise it has a right subtree → replace with successor
+        else if let right = node.right {
+            var successor = right
+            while let l = successor.left { successor = l }
+            node.key = successor.key
+            delete(successor)
+        }
+    }
+
+    /// Walks upward from `n`, fixing heights, balances and performing rotations.
+    private func rebalance(_ n: Node) {
+        setBalance(of: n)
+
+        var node = n
+        if node.balance == -2 {
+            // left heavy
+            if height(of: node.left?.left) >= height(of: node.left?.right) {
+                node = rotateRight(node)
+            } else {
+                node = rotateLeftThenRight(node)
+            }
+        } else if node.balance == 2 {
+            // right heavy
+            if height(of: node.right?.right) >= height(of: node.right?.left) {
+                node = rotateLeft(node)
+            } else {
+                node = rotateRightThenLeft(node)
+            }
+        }
+
+        // continue upwards – or make this node the new root
+        if let parent = node.parent {
+            rebalance(parent)
+        } else {
+            root = node
+        }
+    }
+
+    // ---------- Rotations -------------------------------------------------
+
+    private func rotateLeft(_ a: Node) -> Node {
+        guard let b = a.right else { return a }      // safety guard
+
+        // detach b from a
+        b.parent = a.parent
+        a.right = b.left
+        a.right?.parent = a
+
+        // attach a under b
+        b.left = a
+        a.parent = b
+
+        // reconnect b with the rest of the tree
+        if let p = b.parent {
+            if p.left === a {
+                p.left = b
+            } else {
+                p.right = b
+            }
+        }
+
+        setBalance(of: a, b)
+        return b
+    }
+
+    private func rotateRight(_ a: Node) -> Node {
+        guard let b = a.left else { return a }       // safety guard
+
+        b.parent = a.parent
+        a.left = b.right
+        a.left?.parent = a
+
+        b.right = a
+        a.parent = b
+
+        if let p = b.parent {
+            if p.left === a {
+                p.left = b
+            } else {
+                p.right = b
+            }
+        }
+
+        setBalance(of: a, b)
+        return b
+    }
+
+    private func rotateLeftThenRight(_ n: Node) -> Node {
+        if let left = n.left {
+            n.left = rotateLeft(left)
+        }
+        return rotateRight(n)
+    }
+
+    private func rotateRightThenLeft(_ n: Node) -> Node {
+        if let right = n.right {
+            n.right = rotateRight(right)
+        }
+        return rotateLeft(n)
+    }
+
+    // ---------- Height / Balance helpers ----------------------------------
+
+    /// Height of a node – `-1` for `nil` (matches the Java implementation).
+    private func height(of node: Node?) -> Int {
+        node?.height ?? -1
+    }
+
+    /// Re‑computes stored height of `node`.
+    private func reheight(_ node: Node?) {
+        guard let node = node else { return }
+        node.height = 1 + max(height(of: node.left), height(of: node.right))
+    }
+
+    /// Updates both `height` and `balance` for every supplied node.
+    private func setBalance(of nodes: Node...) {
+        for n in nodes {
+            reheight(n)
+            n.balance = height(of: n.right) - height(of: n.left)
+        }
+    }
+
+    // ---------- Printing ---------------------------------------------------
+
+    private func printBalance(node: Node?) {
+        guard let node = node else { return }
+        printBalance(node: node.left)
+        print("\(node.balance) ", terminator: "")
+        printBalance(node: node.right)
+    }
+}
+
+// MARK: - Demo -------------------------------------------------------------
+
+let tree = AVLTree()
+
+print("Inserting values 1 to 10")
+for i in 1...10 {
+    _ = tree.insert(i)
+}
+
+print("Printing balance: ", terminator: "")
+tree.printBalance()
+print()   // newline