Keep track of a Sorted array using Binary Search

SortedArray.swift

/// Solves the problem https://leetcode.com/problems/count-of-smaller-numbers-after-self/
/// keeping track of elements in sorted order.
/// Ex: A: [4,2,3,1] -> [3,1,1,0]
/// We traverse A in reversed order keeping track of each element in a sorted manner.
/// using the index where the new element was inserted as a counter for the elements after Self.
///
///     0 1 2 3
/// A: [1,3,2,4]
///     i
/// i | sortedArray | result
/// 0      [1]         [0]
/// 1     [1,3]       [0,1]
/// 2    [1,2,3]     [0,1,1]
/// 3   [1,2,3,4]   [0,1,1,3]
/// Return the result array in reversed order

/// Uses binary search to keep track of a sorted array
/// Overall It runs in O(nlogn)
struct SortedArray<T: Comparable> {
    private(set) var data: [T] = []
    
    init(capacity: Int = 0) {
        if capacity > 0 {
            data.reserveCapacity(capacity)
        }
    }
    
    /// Inserts a new element in a sorted order using binary search
    /// Returns the left most index where the new element was inserted
    mutating func insert(_ newElement: T) -> Int {
        if data.isEmpty {
            data.append(newElement)
            return 0
        } else if data.count == 1 {
            if newElement <= data[0] {
                data.insert(newElement, at: 0)
                return 0
            } else {
                data.append(newElement)
                return 1
            }
        } else { // Use binary search to insert new element
            let index = findIndex(newElement)
            data.insert(newElement, at: index)
            return index
        }
    }
    
    
    /// Will find the left most index to insert the new element,
    /// Ex: [1,2,3], inserting 2, it will return the index 1 (first index of 2)
    /// Ex: [2,3,4], inserting 1, it will return 0
    /// Ex: [1,2,3,5], inserting 4, it will return 3
    
    mutating private func findIndex(_ newElement: T) -> Int {
        var low = 0
        var high = data.count - 1
        
        while low <= high {
            let mid = (low + high) / 2
            
            if mid > 0 && mid < data.count && newElement > data[mid - 1] && newElement <= data[mid] {
                return mid
            } else {
                if newElement <= data[mid] {
                    high = mid - 1
                } else {
                    low = mid + 1
                }
            }
        }
        
        return low
    }
}

func countSmallerElementsAfterSelf(_ arr: [Int]) -> [Int] {
    var sortedArray = SortedArray<Int>(capacity: arr.count)
    var result: [Int] = []
    
    for n in arr.reversed() {
        let index = sortedArray.insert(n)
        result.append(index)
    }
    
    return result.reversed()
}