Longest subarray with a zero sum

Question

I'm working through problems on this site.

The problem I'm currently doing is:

Given an array having both positive and negative integers, your task is to complete the function maxLen which returns the length of maximum subarray with 0 sum. The function takes two arguments an array A and n where n is the size of the array A.

However the code is too slow to pass so I can't check if the solution is working correctly or not - (it works fine for the single test case that's provided)

Can anyone suggest how to improve the performance of this code?

from itertools import combinations

def maxLen(n, arr):
    for i in range(n, 0, -1):
        for combo in combinations(arr, i):
            if sum(combo) == 0:
                return len(combo)
    else:
        return 0

Answer 1

(EDIT: Go look at @RE60K's solution before mine. His is faster and easier to understand.)

Unfortunately, I don't think this code will work, efficiency aside.

for combo in combinations(arr, i):

Here, you're looping through permutations, but the challenge asks for subarrays.

Instead, you should do something like:

def maxLen(n, arr):
    for subarray_length in range(n, 0, -1):
        for subarray_start in range(len(arr) - subarray_length + 1):
            subarray = arr[subarray_start : subarray_start + subarray_length]
            if sum(subarray) == 0: return len(subarray)
    return 0

Which works, but isn't terribly quick. If sum is O(n), then this is O(n^3).

I was able to come up with a quicker algorithm (O(n^2)) which seems to be consistently faster on my timeit tests (which may not be perfect).

To show the idea of how this works, I'll use an example.

Consider if the input array were [1, -4, 4, -3, -1, 8, -3].

We'd start by trying the whole array, which sums to -4. Then we'd try all arrays one size down, then two, etc.

Let's skip to trying subarray size 3, since that's where the answer lies. First we try [1, -4, 4] which sums to 1. Normally, we'd then try the next subarray, [-4, 4, -3], but we can take a shortcut here. These two subarrays only differ by two items: the first item of the first, and the last item of the second. As such, we can instead keep a running sum, and only worry ourselves with those two key items.

So, instead, we go like this: first length 3 subarray is [1, -4, 4], which sums to 1. The next item after this subarray is -3. We add this to the running sum, and remove 1, (+= -3, -= 1), and then consider our array to have moved to the right one. So, the next item we consider is the one after -3, namely, -1. And so we continue and find that our running sum does hit 0, and so return 3.

def max_len_2(arr_len, arr):
    for sub_len in range(arr_len, 0, -1):  # For all subarray lengths
        arr_sum = sum(arr[0:sub_len])  # Get the initial sum
        if arr_sum == 0: return sub_len  # Return if it works
        for gain_index in range(sub_len, arr_len):  # Otherwise, loop through all the indices after the array
            arr_sum -= arr[gain_index - sub_len]  # Drop first item in array
            arr_sum += arr[gain_index]  # Gain first item after array
            if arr_sum == 0: return sub_len  # If sum works, return
    return 0

This is more efficient, awesome! But we can take it one step further. What we just did is we used information we already know about a row in order to have to make less calculations. Essentially, we're not wasting any information.

And we can do this between rows, as well, in order to be even more efficient (although still O(n^2), but whatever...).

I'll explain this algorithm with the same example, [1, -4, 4, -3, -1, 8, -3].

We start by noting that the sum of this is -4. We'll call this sum 0 0. If it were 0, we'd be done, but it isn't.

Now, we start with the largest subarrays, length 6. We note that the first subarray is the same as the whole array, except for the last element; so, the sum is sum 0 0 minus the last element, which is -3. The sum is -4 - -3 = -1. This is sum 1 0. The second length-6 subarray is the same as the whole array, except for the first element, so we know the sum is sum 0 0 - first element = -4 - 1 = -5.

Now, onto length 5. The sum of the first length 5 subarray (sum 2 0) is the same as the first length 6 subarray (sum 1 0), except for the first element...

See the pattern? With each new subarray length n, sum n k = sum (n-1) k - last element of array (n-1) k. The only exception here is that there is 1 length-7 sub, 2 length-6, 3 length-5, etc.; each time we gain a subarray, and this new sub isn't accounted for with this algorithm. Instead, we know that the sum of the last sub sum n k is the same as sum (n-1) k - first element of array (n-1) k Note the difference: it's the first element we're taking away, not the last.

def max_len_dynamic_programming(arr_len, arr):
    dp = []  # This is our data storage

    # We initialize it with the sum of the array
    s = sum(arr)
    if s == 0: return arr_len
    dp.append([s])

    for sub_len in range(arr_len - 1, 0, -1):  # Iterate over sub lengths, decreasing...
        row = []  # We'll construct the new row
        prev_row = dp[-1]  # With info from the old row
        prev_row_len = arr_len - sub_len

        for i in range(prev_row_len):
            # This is where the magic happens.
            # cell is `sum n k`, and
            # prev_row[i] is `sum (n-1) k`,
            # and arr[i + sub_len] is
            # the last item of `array (n-1) k`
            cell = prev_row[i] - arr[i + sub_len]
            if cell == 0: return sub_len
            row.append(cell)

        # Deal with the last element seperately
        last = prev_row[-1] - arr[~sub_len]  # -sub_len - 1
        if last == 0: return sub_len
        row.append(last)

        dp.append(row)

    return 0

If you want a faster*, harder-to-understand version, make these swaps in the code:

Swap

row = []
prev_row = dp[-1]
prev_row_len = arr_len - sub_len

for

prev_row = dp[-1]
prev_row_len = arr_len - sub_len
row_len = prev_row_len + 1
row = [0] * row_len

& swap row.append(cell) for row[i] = cell

& swap row.append(last) for row[-1] = last

*Most of the time, in my tests, at least. The swaps change this from using all dynamic arrays to preinitialized arrays for each row, and dynamic for the grid. Going full-preinitialized-arrays will be even faster some of the time. Not sure how to tell when it'll be faster and when it'll be slower.

Answer 2

I found @Quelklef's \$O(n^2)\$ solution quite complicated and I think this would be a better one, with \$O(n)\$ storage for partial sums, i.e. \$S_k=\sum_{i=0}^ka_i\$ and sum of a sub-array from index \$a\$ to \$b\$ is \$S_b-S_{a-1}\$

def maxLen(n, arr):
    sm = [0] # sm[1] = a[0], sm[2] = a[0]+a[1], ...
    for i in range(0, n): # sm[k] = sm[k-1] + arr[k-1]
        sm.append(sm[-1] + arr[i])
    for j in range(n, 0, -1): # with length j, i.e. arr[i]+arr[i+1]+...arr[i+j-1], i.e. sm[i+j]-sm[i]
        for i in range(0, n-j+1): # we are taking the array from index i    
            if sm[i+j]-sm[i] == 0:
                return j
    return 0