IntervalTree.swift

import Foundation

// MARK: - Interval
struct Interval {
    var low: Int
    var high: Int
    
    func overlaps(_ otherInterval: Interval) -> Bool {
        return contains(otherInterval.low) || contains(otherInterval.high)
    }
    
    private func contains(_ int: Int) -> Bool {
        return low <= int && int <= high
    }
}

extension Interval: Equatable {
    static func == (lhs: Self, rhs: Self) -> Bool {
        return lhs.low == rhs.low && lhs.high == rhs.high
    }
}

extension Interval: CustomStringConvertible {
    var description: String {
        return "[\(low), \(high)]"
    }
}

// MARK: - ITNode
class ITNode {
    // Properties
    let interval: Interval
    fileprivate(set) var max: Int
    
    // Children (Subtrees)
    var left: IntervalTree?
    var right: IntervalTree?
    
    init(interval: Interval) {
        self.interval = interval
        self.max = interval.high
    }
    
    var isLeaf: Bool {
        return left == nil && right == nil
    }
    
}

extension ITNode: CustomStringConvertible {
    var description: String {
        return "\(interval.description) \(max)"
    }
}


// MARK: - IntervalTree
class IntervalTree {
    
    var root: ITNode
    fileprivate(set) var height: Int
    init(_ root: ITNode) {
        self.root = root
        self.height = 1
    }
    
    /// Inserts an interval into the `IntervalTree`. If root is `nil` a new root is created and the `interval` parameter will be the root's interval.
    /// **Repeated Intervals are ignored.**
    /// - Parameter interval: Interval to be inserted.
    func insert(_ interval: Interval) {
        // Ignore already existing intervals
        if contains(interval) { return }
        
        // Update max of the root
        root.max = max(root.max, interval.high)
        
        // Lower than root we go left
        if interval.low <= root.interval.low {
            if root.left == nil {
                root.left = IntervalTree(ITNode(interval: interval))
            } else {
                root.left?.insert(interval)
            }
        }
        // Higher than root we go right
        else {
            if root.right == nil {
                root.right = IntervalTree(ITNode(interval: interval))
            } else {
                root.right?.insert(interval)
            }
        }
        
        // Applying balancing mechanism
    }
    
    /// Returns the first overlaped interval with the provided interval using DFS. Returns `nil` if not found.
    /// - Parameter interval: The `targetInterval` for search.
    func overlaps(_ interval: Interval) -> Interval? {
        if root.interval.overlaps(interval) { return root.interval }
        return root.left?.overlaps(interval) ?? root.right?.overlaps(interval)
    }
    
    /// Return `Bool` whether the provided interval does exists in the tree or not.
    /// - Parameter interval: The `targetInterval` for search
    func contains(_ interval: Interval) -> Bool {
        if root.interval == interval { return true }
        
        if interval.low <= root.interval.low {
            return root.left?.contains(interval) ?? false
        } else {
            return root.right?.contains(interval) ?? false
        }
    }
    
    // MARK: - Balancing helpers
    func rotateRight(_ subtree: IntervalTree) -> IntervalTree {
        // Get references
        let leftChild = subtree.root.left
        let rightChildOfLeftChild = leftChild?.root.right
        
        // Do the rotation
        leftChild?.root.right = subtree
        subtree.root.left = rightChildOfLeftChild
        
        return leftChild!
    }
    
    func rotateLeft(_ subtree: IntervalTree) -> IntervalTree {
        // Get references
        let rightChild = subtree.root.right
        let leftChildOfRightChild = rightChild?.root.left
        
        // Do the rotation
        rightChild?.root.left = subtree
        subtree.root.right = leftChildOfRightChild
        
        return rightChild!
    }
    
    // MARK: - Graph
    private var margin: Int {
        return 13
    }
    
    func printGraph() {
        print(graph2D(root, space: 0))
    }
    
    private func graph2D(_ root: ITNode?, space: Int) -> String {
        guard let root = root else { return "" }
        
        var graph = ""
        graph += graph2D(root.right?.root, space: space + margin)
        for _ in 0..<space { graph += " " }
        graph += "\(root.description)\n"
        graph += graph2D(root.left?.root, space: space + margin)
        
        return graph
    }
}

extension IntervalTree: CustomStringConvertible {
    var description: String {
        var inOrderTraversal = ""
        inOrderTraversal += root.left?.description ?? ""
        inOrderTraversal += root.description + "\n"
        inOrderTraversal += root.right?.description ?? ""
        return inOrderTraversal
    }
}


// MARK: - Testing
print("================================================")
print("Test Insertion")
print("================================================")
let tree = IntervalTree(ITNode(interval: Interval(low: 10, high: 20)))
tree.printGraph()
print("================================================")
tree.insert(Interval(low: 5, high: 15))
tree.printGraph()
print("================================================")
tree.insert(Interval(low: 15, high: 45))
tree.printGraph()
print("================================================")
tree.insert(Interval(low: 12, high: 25))
tree.printGraph()
print("================================================")
tree.insert(Interval(low: 17, high: 23))
tree.printGraph()
print("================================================")
tree.insert(Interval(low: 16, high: 23))
tree.printGraph()


// Test Contains
print("================================================")
print("Test Contains")
print("================================================")
print("Contains [5, 15]?", "\t\t", tree.contains(Interval(low: 5, high: 15)))
print("Contains [0, 1]?", "\t\t", tree.contains(Interval(low: 0, high: 1)))


// Test Overlaps
print("================================================")
print("Test Overlaps")
print("================================================")
print("Overlaps [40, 45]?", "\t\t", tree.overlaps(Interval(low: 40, high: 45)) ?? "No overlapping found")
print("Overlaps [0, 1]?", "\t\t", tree.overlaps(Interval(low: 0, high: 1)) ?? "No overlapping found")


// Test rotation at root
var rotatedTree = tree
print("================================================")
print("Test Rotation")
print("================================================")
print("Original Tree")
rotatedTree.printGraph()
print("================================================")
print("Rotated Left")
rotatedTree = rotatedTree.rotateLeft(rotatedTree)
rotatedTree.printGraph()
print("================================================")
print("Rotated Right")
rotatedTree = rotatedTree.rotateRight(rotatedTree)
rotatedTree.printGraph()
print("================================================")
print("Rotated Right Again")
rotatedTree = rotatedTree.rotateRight(rotatedTree)
rotatedTree.printGraph()
print("================================================")