RangeTree.swift

class RangeTree<T> { 
    
    private var value: [T]
    private(set) var range: Range<Int>
    private var leftLeaf: RangeTree<T>?
    private var rightLeaf: RangeTree<T>?
    
    convenience init(array: [T]) {
        self.init(array: array, range: 0 ..< array.count)
    }
    
    init(array: [T], range: Range<Int>) {
        self.range = range
        
        if range.lowerBound == range.upperBound {
            if array.indices.contains(range.upperBound) {
                self.value = [array[range.lowerBound]]
            } else {
                self.value = []
            }
        } else {
            let middle = (range.lowerBound + range.upperBound) / 2
            leftLeaf = RangeTree(array: array, range: range.lowerBound ..< middle)
            rightLeaf = RangeTree(array: array, range: (middle + 1)..<range.upperBound)
            
            value = leftLeaf!.value + rightLeaf!.value
        }
    }
    
    func token(by range: NSRange) -> [T] {
        return self.token(by: range.lowerBound..<range.upperBound)
    }
    
    func token(by range: Range<Int>) -> [T] {
        if self.range.lowerBound == range.lowerBound && self.range.upperBound <= range.upperBound {
            return self.value
        }
        
        guard let leftChild = self.leftLeaf else { return [] }
        guard let rightChild = self.rightLeaf else { return [] }
        
        if leftChild.range.upperBound < range.lowerBound {
            return rightChild.token(by: range)
        } else if rightChild.range.lowerBound > range.upperBound {
            return leftChild.token(by: range)
        } else {
            let leftResult = leftChild.token(by: range)
            let rightResult = rightChild.token(by: range)
            return leftResult + rightResult
        }
    }
    
}