How to get O(1) time complexity for Array.remove(at:) in Swift

Viewed 1601

I know usually it takes O(n) time to remove an item from a Swift Array by using remove(at:).

Let us say I have some code like this:

// refArray has m elements
// myArray has n elements

for refItem in refArray {
    for (index, item) in myArray.enumerated() {
        if item.id == refItem.id {
            <Do something>
            myArray.remove(at: index)
            break
        }
    }
}

for myArray.remove(at: index), even I already know the index of the item to be deleted, but if I use remove(at:), it still takes O(n) time. So the whole process could take O(n³) time.

Is there any way to make it O(n²) but still removing an item from myArray when item.id == refItem.id?

3 Answers

If you want to remove m items from an array of n items, yes, you could refactor that to achieve O(m*n) time complexity. But whenever you see nested for loops (or nested functions that are tantamount to the same thing), you should consider whether there are structures you could build that will cut that O(m*n) down to O(m+n), which is dramatically faster.

In this case, we do have such an opportunity. So, we first build a hashed structure of the values for which we will be searching. Then we can iterate through the array once looking for records to be removed, each with O(1) complexity. And as discussed elsewhere, rather than having your own for loop calling remove(at:), we can employ methods like removeAll(where:) will do this looping and removing all in one step of O(n) complexity.

extension Array where Element: Hashable {
    mutating func remove(values: [Element]) {
        let set = Set(values)
        removeAll { set.contains($0) }
    }
}

And then

var array = [0, 1, 2, 3, 4, 5, 6]
var values = [1, 3, 5]
array.remove(values: values)
print(array)                       // [0, 2, 4, 6]

There are lots of variations on the theme, but the idea is generally the same: If you can, build a structure where values to be removed can be identified in O(1) time. Then you can iterate through the array only once, e.g. using removeAll. Then net result is an overall O(m+n) time complexity, which is far faster than either your original algorithm or the contemplated O(m*n) rendition.

Note, patterns like this work if (a) the elements in the array are Hashable; and (b) we are OK increasing the space complexity of the algorithm. We are, in short, sacrificing space complexity for improvements in the time complexity.

If myArray does not contain items with duplicate IDs, you can use removeAll(where:) instead of your double for loops:

myArray.removeAll(where: { item in 
    refArray.contains { $0.id == item.id }
})

This is O(n*m), where n and m are the size of myArray and refArray respectively.

Otherwise, you can do what starboy_b says in the comments, collect all the indices and remove them in one go with remove(atOffsets:):

var indicesToRemove = IndexSet()
for refItem in refArray {
    for (index, item) in myArray.enumerated() {
        if i == j, !indicesToRemove.contains(index) {
            indicesToRemove.insert(index)
            break
        }
    }
}
myArray.remove(atOffsets: indicesToRemove)

This is also O(n*m).

If order isn't important, you can use swapAt(indexOfElementToRemove, lastIndex) and then removeLast(). This achieves O(1) time. However, a Set may be a better type to use.

Related