Swift - efficient array searching

Question

I am trying to find a way to search an Int array (all numbers are > 0) and identify sequences where the sum of each element is a specific number. One element can be in more than one sequence.

This works:

let input = [2, 52, 23, 52, 28, 47, 13, 15]
let target = 75

var result: [[Int]] = []

let count = input.count
var start = 0

while start < count {
    for index in start..<count {
        let temp = Array(input[start...index])
        let sum = temp.reduce(0, +)
        
        if sum == target {
            result.append(temp)
            break
        }
        
        if sum > target {
            break
        }
    }
    
    start += 1
}


print(result) // [[52, 23], [23, 52], [28, 47], [47, 13, 15]]

However, when the array gets longer and the target higher, it obviously slows down. For instance an array of 10000 elements and a target of 400 takes about 2 seconds. The above example takes 0.008 seconds.

Is there anyway I can improve on my code so that it doesn't slow down so much for larger arrays and targets?

Answer 1

Instead of putting everything into the “main” unit it is better to put the search algorithm in a dedicated function:

func subarrays(in input:[Int], withSum target: Int) -> [[Int]] {
    // ...
}

print(subarrays(in: [2, 52, 23, 52, 28, 47, 13, 15], withSum: 75))

That is clearer and allows to add test cases easily.

You use a while loop for iterating over the start index, and a for loop for iterating over the end index. I would use a for loop for both.

If the lower index is called start end the upper index is better called end (alternatively: startIndex, endIndex or lowIndex, highIndex).

Since all input numbers are positive, a target <= 0 can never be achieved, so one can check that case first.

You add the input numbers between start and end index with reduce(), which can be done directly on the array slice. This saves some array creations (which copy the elements).

The code would then look like this:

func subarrays(in input:[Int], withSum target: Int) -> [[Int]] {
    if target <= 0 {
        return []
    }
    
    var result: [[Int]] = []
    let count = input.count
    for start in 0..<count {
        for end in start..<count {
            let sum = input[start...end].reduce(0, +)
            if sum == target {
                result.append(Array(input[start...end]))
                break
            }
            if sum > target {
                break
            }
        }
    }
    return result
}

Instead of calculating the sum over every slice one can update the sum when the end index is incremented. This allows to get rid of the reduce() call and reduces (!) the complexity from \$ O(N^3) \$ to \$ O(N^2) \$ (where \$ N \$ is the length of the input array):

func subarrays(in input:[Int], withSum target: Int) -> [[Int]] {
    if target <= 0 {
        return []
    }

    var result: [[Int]] = []
    let count = input.count
    for start in 0..<count {
        var sum = 0
        for end in start..<count {
            sum += input[end]
            if sum == target {
                result.append(Array(input[start...end]))
                break
            }
            if sum > target {
                break
            }
        }
    }
    return result
}

But there is a better algorithm: Instead of two nested loops it suffices to increment the end index if the current sum is too small, and to increment the start index if the current sum is too large. The current sum is always kept up to date by incrementing or decrementing it accordingly:

func subarrays(in input:[Int], withSum target: Int) -> [[Int]] {
    var result: [[Int]] = []

    let count = input.count
    var start = 0
    var end = 0
    var sum = 0
    // Loop invariants: 
    //    0 <= start <= end <= count
    //    sum = input[start] + ... + input[end - 1]
    while start < count {
        while end < count {
            sum += input[end] ; end += 1
            if sum == target {
                result.append(Array(input[start..<end]))
                break
            } else if sum > target {
                break
            }
        }
        while start < end {
            sum -= input[start] ; start += 1
            if sum == target {
                result.append(Array(input[start..<end]))
                break
            } else if sum < target {
                break
            }
        }
    }
    return result
}

Documenting the loop invariants help to understand the logic of the nested while loops.

This algorithm has \$ O(N) \$ complexity, and should be considerably faster for large arrays.