AVLTree.swift

//
//  AVLTree.swift
//  AVLTree
//
//  Swift port of immutable AVLTree I implemented with Stephan Partzsch
//
//  Original ObjC implementation: https://github.com/StephanPartzsch/AVLTree
//
//  Copyright (c) 2014 Maxim Zaks. All rights reserved.
//

func ||<T>(optional : Optional<T>, defaultValue : T) -> T {
    if let value = optional {
        return value
    }
    return defaultValue
}

public func +<T>(node : AVLNode<T>, newValue : T) -> AVLNode<T> {
    var newLeft : AVLNode<T>? = node.left
    var newRight : AVLNode<T>? = node.right
    if(newValue < node.value) {
        if let left = node.left {
            newLeft = left + newValue
        } else {
            newLeft = AVLNode(newValue)
        }
    } else if(newValue > node.value) {
        if let right = node.right {
            newRight = right + newValue
        } else {
            newRight = AVLNode(newValue)
        }
    } else {
        return node
    }
    
    let newRoot = AVLNode(value: node.value, left: newLeft, right: newRight)
    
    return newRoot.fixBalance()
}

public func-<T>(node: AVLNode<T>?, value : T) -> AVLNode<T>? {
    return node?.remove(value).result
}



public final class AVLNode<T : Comparable> {
    typealias Element = T
    let left : AVLNode<Element>?
    let right : AVLNode<Element>?
    public let count : UInt
    let depth : UInt
    let balance : Int
    let value : Element!
    
    public convenience init(_ value : Element){
        self.init(value: value, left: nil, right: nil)
    }
    
    public func contains(value : Element) -> Bool{
        if self.value == value {
            return true
        }
        if left?.contains(value) == true {
            return true
        }
        if right?.contains(value) == true {
            return true
        }
        
        return false
    }
    
    init(value : Element, left: AVLNode<Element>?, right: AVLNode<Element>?){
        self.value = value
        self.left = left
        self.right = right
        self.count = 1 + (left?.count || 0) + (right?.count || 0)
        self.depth = 1 + max((left?.depth || 0), (right?.depth || 0))
        self.balance = (left?.depth || 0) - (right?.depth || 0)
    }
    
    func fixBalance() -> AVLNode<Element> {
        if abs(balance) < 2 {
            return self
        }
        
        if (balance == 2)
        {
            let leftBalance = self.left?.balance || 0
            
            if (leftBalance == 1 || leftBalance == 0)
            {
                //Easy case:
                return rotateToRight()
            }
            
            if (leftBalance == -1)
            {
                //Rotate Left to left
                let newLeft = left!.rotateToLeft()
                let newRoot = AVLNode(value: value, left: newLeft, right: right)
                
                return newRoot.rotateToRight()
            }
            
            fatalError("LeftNode too unbalanced")
        }
        
        if (balance == -2)
        {
            let rightBalance = right?.balance || 0
            
            if (rightBalance == -1 || rightBalance == 0)
            {
                //Easy case:
                return rotateToLeft()
            }
            
            if (rightBalance == 1)
            {
                //Rotate right to right
                let newRight = right!.rotateToRight()
                let newRoot = AVLNode(value: value, left: left, right: newRight)
                
                return newRoot.rotateToLeft()
            }
            
            fatalError("RightNode too unbalanced")
        }
        
        fatalError("Tree too unbalanced")
        
        return self
    }
    
    func remove(value : Element) -> (result: AVLNode<Element>?, foundFlag :Bool){
        if value < self.value {
            
            let removeResult = left?.remove(value)
            
            if removeResult == nil || removeResult!.foundFlag == false {
                // Not found, so nothing changed
                return (self, false)
            }
            
            let newRoot = AVLNode(value: self.value, left: removeResult!.result, right: right).fixBalance()
            
            return (newRoot, true)
        }
        
        if value > self.value {
            let removeResult = right?.remove(value)
            
            if removeResult == nil || removeResult!.foundFlag == false {
                // Not found, so nothing changed
                return (self, false)
            }
            
            let newRoot = AVLNode(value: self.value, left: left, right: removeResult!.result)
            
            return (newRoot, true)
        }
        
        //found it
        return (removeRoot(), true)
    }
    
    func removeMin()-> (min : AVLNode<Element>, result : AVLNode<Element>?){
        if left == nil {
            //We are the minimum:
            return (self, right)
        } else {
            //Go down:
            let (min, newLeft) = left!.removeMin()
            let newRoot = AVLNode(value: value, left: newLeft, right: right)
            
            return (min, newRoot.fixBalance())
        }
    }
    
    func removeMax()-> (max : AVLNode<Element>, result : AVLNode<Element>?){
        if right == nil {
            //We are the max:
            return (self, left)
        } else {
            //Go down:
            let (max, newRight) = right!.removeMax()
            let newRoot = AVLNode(value: value, left: left, right: newRight)
            
            return (max, newRoot.fixBalance())
        }
    }
    
    func removeRoot() -> AVLNode<Element>? {
        if left == nil {
            return right
        }
        
        if right == nil {
            return left
        }
        
        //Neither are empty:
        if left!.count < right!.count {
            // LeftNode has fewer, so promote from RightNode to minimize depth
            let (min, newRight) = right!.removeMin()
            let newRoot = AVLNode(value: min.value, left: left, right: newRight)
            
            return newRoot.fixBalance()
        }
        else
        {
            let (max, newLeft) = left!.removeMax()
            let newRoot = AVLNode(value: max.value, left: newLeft, right: right)
            
            return newRoot.fixBalance()
        }
    }

    func rotateToRight() -> AVLNode<Element> {
        let newRight = AVLNode(value: value, left: left!.right, right: right)
        return AVLNode(value: left!.value, left: left!.left, right: newRight)
    }
    
    func rotateToLeft() -> AVLNode<Element> {
        let newLeft = AVLNode(value: value, left: left, right: right!.left)
        return AVLNode(value: right!.value, left: newLeft, right: right!.right)
    }
}

extension AVLNode : Printable, DebugPrintable {
    public var description : String {
        get {
            let empty = "_"
            return "(\(value) \(left?.description || empty) \(right?.description || empty))"
        }
    }
    
    public var debugDescription : String {
        get {
            return self.description
        }
    }
}


//  SequenceType implementation is based on following blog post:
//  http://www.spazcosoft.com/posts/data-structures-in-swift-part-1-wow-thats-an-unstable-compiler/

extension AVLNode : SequenceType {
    public func generate() -> GeneratorOf<Element> {
        
        var stack : [AVLNode<Element>] = []
        var current : AVLNode<Element>? = self
        
        return GeneratorOf<Element> {
            while(stack.count != 0 || current != nil){
                if(current != nil) {
                    stack.append(current!)
                    current = current!.left
                } else {
                    let retval = stack.removeLast();
                    current = retval.right
                    return retval.value
                }
            }
            
            return nil
        }
    }
}

struct ReversedSequenceOfAvlNode<T : Comparable> : SequenceType {
    let node : AVLNode<T>
    func generate() -> GeneratorOf<T> {
        
        var stack : [AVLNode<T>] = []
        var current : AVLNode<T>? = node
        
        return GeneratorOf<T> {
            while(stack.count != 0 || current != nil){
                if(current != nil) {
                    stack.append(current!)
                    current = current!.right
                } else {
                    let retval = stack.removeLast();
                    current = retval.left
                    return retval.value
                }
            }
            
            return nil
        }
    }
}

extension AVLNode {
    public var reversed : SequenceOf<Element> {
        get {
            return SequenceOf(ReversedSequenceOfAvlNode(node: self))
        }
    }
}

AVLTreeTests.swift

//
//  AVLTreeTests.swift
//  AVLTreeTests
//
//  Created by Maxim Zaks on 25.12.14.
//  Copyright (c) 2014 Maxim Zaks. All rights reserved.
//

import UIKit
import XCTest
import AVLTree

class AVLTreeTests: XCTestCase {
    
    func testTreeAdition() {
        var tree = AVLNode("D") + "A"
        println("1-------------- \(tree.description) -------------")
        XCTAssertEqual(tree.description, "(D (A _ _) _)", "expected tree")
        
        tree = tree + "H"
        println("2-------------- \(tree.description) -------------")
        XCTAssertEqual(tree.description, "(D (A _ _) (H _ _))", "expected tree")
        
        tree = tree + "C"
        println("3-------------- \(tree.description) -------------")
        XCTAssertEqual(tree.description, "(D (A _ (C _ _)) (H _ _))", "expected tree")
        
        tree = tree + "A"
        println("4-------------- \(tree.description) -------------")
        XCTAssertEqual(tree.description, "(D (A _ (C _ _)) (H _ _))", "expected tree")
        
        tree = tree + "U"
        println("5-------------- \(tree.description) -------------")
        XCTAssertEqual(tree.description, "(D (A _ (C _ _)) (H _ (U _ _)))", "expected tree")
        
        tree = tree + "E"
        println("6-------------- \(tree.description) -------------")
        XCTAssertEqual(tree.description, "(D (A _ (C _ _)) (H (E _ _) (U _ _)))", "expected tree")
    }
    
    func testBalance() {
        var tree = AVLNode("A") + "B" + "C"
        println("-------------- \(tree.description) -------------")
        XCTAssertEqual(tree.description, "(B (A _ _) (C _ _))", "balance should be fixed")
    }
    
    func testRemove(){
        var tree = AVLNode("A") + "B" + "C"
        XCTAssertEqual(tree.description, "(B (A _ _) (C _ _))", "balance should be fixed")
        
        let tree1 = tree - "A"
        XCTAssertEqual(tree1!.description, "(B _ (C _ _))", "removed A from A B C")
        
        let tree2 = tree - "C"
        XCTAssertEqual(tree2!.description, "(B (A _ _) _)", "removed C from A B C")
        
        let tree3 = tree - "B"
        XCTAssertEqual(tree3!.description, "(A _ (C _ _))", "removed B from A B C")
        
        let tree4 = tree - "F"
        XCTAssertEqual(tree4!.description, "(B (A _ _) (C _ _))", "removed F from A B C")
        
        let tree5 = tree - "A" - "B" - "C" - "F"
        XCTAssertNil(tree5, "removed all values")
        
    }
    
    func testSequene(){
        let tree = AVLNode(6) + 2 + 9 + 4 + 5
        let array : [Int] = map(tree, {$0})
        XCTAssertEqual(array, [2, 4, 5, 6, 9], "elements should be iterated through in the right order")
    }
    
    func testReverseSequene(){
        let tree = AVLNode(6) + 2 + 9 + 4 + 5
        let array : [Int] = map(tree.reversed, {$0})
        XCTAssertEqual(array, [9, 6, 5, 4, 2], "elements should be iterated through in the right order")
    }
    
    func testContains(){
        let tree = AVLNode(5) + 1 + 6
        XCTAssert(tree.contains(5))
        XCTAssert(tree.contains(1))
        XCTAssert(tree.contains(6))
        XCTAssert(!tree.contains(2))
        XCTAssert(!tree.contains(-2))
        XCTAssert(!tree.contains(20))
    }
    
    func testCount(){
        let tree = AVLNode(5) + 1 + 6
        XCTAssert(tree.count == 3)
    }
}

This gist is released under the MIT licence. Feel free to do what ever you want with it!