Augmented Interval Search tree

AugmentedIntevalTree.swift

//
// Created by @fewlinesofcode on 9/6/18.
// Copyright (c) 2018 Oleksandr Glagoliev. All rights reserved.
//

public class Interval<T: Comparable> {
    private (set) var start: T
    private (set) var end: T
    var max: T
    
    var left: Interval<T>?
    var right: Interval<T>?
    
    init(start: T, end: T) {
        precondition(start <= end)
        
        self.start = start
        self.end = end
        self.max = end
    }
}

extension Interval: Comparable {
    public static func ==(lhs: Interval<T>, rhs: Interval<T>) -> Bool {
        return lhs.start == rhs.start && lhs.end == rhs.end
    }
    public static func <(lhs: Interval<T>, rhs: Interval<T>) -> Bool {
        return lhs.start < rhs.start
    }
    public static func <=(lhs: Interval<T>, rhs: Interval<T>) -> Bool {
        return lhs < rhs || lhs == rhs
    }
    public static func >=(lhs: Interval<T>, rhs: Interval<T>) -> Bool {
        return lhs > rhs || lhs == rhs
    }
    public static func >(lhs: Interval<T>, rhs: Interval<T>) -> Bool {
        return lhs.start > rhs.start
    }
}

extension Interval: CustomStringConvertible {
    public var description: String {
        return "(\(start),\(end)):{\(max)}"
    }
}

public struct AugmentedIntervalTree<T: Comparable> {
    private (set) var root: Interval<T>? = nil
    
    private func insert(_ node: Interval<T>?, _ newNode: Interval<T>) -> Interval<T> {
        guard let tmp = node else {
            return newNode
        }
        
        if newNode.end > tmp.max {
            tmp.max = newNode.end
        }
        
        if (tmp < newNode) {
            if tmp.right == nil {
                tmp.right = newNode
            } else {
                _ = insert(tmp.right, newNode)
            }
        } else {
            if tmp.left == nil {
                tmp.left = newNode
            } else {
                _ = insert(tmp.left, newNode)
            }
        }
        return tmp
    }
    
    func overlaps(acc: inout [Interval<T>], node: Interval<T>?, _ interval: Interval<T>) {
        guard let tmp = node else {
            return
        }
        
        if !((tmp.start > interval.end) || (tmp.end < interval.start)) {
            acc.append(tmp)
        }
        
        if let l = tmp.left, l.max >= interval.start {
            overlaps(acc: &acc, node: l, interval)
        }
        if let r = tmp.right, r.start <= interval.end {
            overlaps(acc: &acc, node: r, interval)
        }
    }
    
    private func printTree(_ tree: Interval<T>? = nil) {
        if tree == nil {
            print("The tree is empty!")
            return
        }
        
        if let left = tree!.left {
            printTree(left)
        }
        
        print(tree!)
        
        if let right = tree!.right {
            printTree(right)
        }
    }
    
    // MARK: Public
    
    public mutating func insert(_ node: Interval<T>) {
        root = insert(root, node)
    }
    
    public func printTree() {
        printTree(root)
    }
    
    public func overlaps(with interval: Interval<T>) -> [Interval<T>]{
        var res = [Interval<T>]()
        overlaps(acc: &res, node: root, interval)
        return res
    }
}

print("--- Tree ---")
var ait = AugmentedIntervalTree<Int>()
ait.insert(Interval<Int>(start: 5, end: 10))
ait.insert(Interval<Int>(start: 15, end: 25))
ait.insert(Interval<Int>(start: 1, end: 12))
ait.insert(Interval<Int>(start: 8, end: 16))
ait.insert(Interval<Int>(start: 14, end: 20))
ait.insert(Interval<Int>(start: 18, end: 21))
ait.insert(Interval<Int>(start: 2, end: 8))
ait.printTree()

print("--- Intersection ---")
let queries = [Interval<Int>(start: 8, end: 10), Interval<Int>(start: 20, end: 22)]

var overlaps = [Interval<Int>]()

for query in queries {
    print("Intersections with \(query): \(ait.overlaps(with:query))")
}
print(overlaps)

Great! @fewlinesofcode thanks for looking into it