Subarray Sum Equals K

Question

I've written a solution to the following leetcode problem:

Given an array of integers and an integer k, you need to find the total number of continuous subarrays whose sum equals to k.

Example 1: Input: nums = [1,1,1], k = 2 Output: 2

Note: The length of the array is in range [1, 20,000]. The range of numbers in the array is [-1000, 1000] and the range of the integer k is [-1e7, 1e7].

Solution:

class Solution:
    def subarraySum(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        T = ['inf' for _ in range(len(nums))]
        count = 0
        for i in range(len(nums)):
            for j in range(i,len(nums)):
                if j == i:
                    T[i] = nums[i]
                    if T[i] == k:
                        count +=1
                else:
                    currSum = T[j-1] + nums[j]
                    T[j] = currSum
                    if currSum == k:
                        count +=1
        return count

The solution passes 58/80 test cases, but for some reason, the solution is returning a Time Limit Exceeded exception on an input array with hundreds of elements. I implemented a dynamic programming solution using a 1d array to memoize previous sums so the algorithm should be efficient. Is there any way to optimize this?

Answer 1

These nested loops

    for i in range(len(nums)):
        for j in range(i,len(nums)):

undoubtedly make the time complexity quadratic. This is the reason for TLE. There is no point to optimize this code; you need to rethink the approach. I don't want to spell out the linear algorithm, and deprive you the fun of finding it. Hint: think of the sliding window.

---

A condition j == i inside the inner loop happens exactly once, and we perfectly know when. Lift it out:

    for i in range(len(nums)):
        T[i] = nums[i]
        if T[i] == k:
            count +=1
        for j in range(i + 1, len(nums)):
            currSum = T[j-1] + nums[j]
            T[j] = currSum
            if currSum == k:
                count +=1

Answer 2

Optimization for code:

1. calculate len(nums) once outside for loops
2. if you use Python2 - use xrange instead of range

Answer 3

You claim you've implemented a dynamic programming solution, but I'm afraid your attempt has actually made your code slower.

Your code:

    T = ['inf' for _ in range(len(nums))]
    count = 0
    for i in range(len(nums)):
        for j in range(i,len(nums)):
            if j == i:
                T[i] = nums[i]
                if T[i] == k:
                    count +=1
            else:
                currSum = T[j-1] + nums[j]
                T[j] = currSum
                if currSum == k:
                    count +=1

For the first change, let's replace the T[i] and nums[i] indices in the j == i case with [j]. Obviously, the code should be equivalent.

    T = ['inf' for _ in range(len(nums))]
    count = 0
    for i in range(len(nums)):
        for j in range(i,len(nums)):
            if j == i:
                T[j] = nums[j]
                if T[j] == k:
                    count +=1
            else:
                currSum = T[j-1] + nums[j]
                T[j] = currSum
                if currSum == k:
                    count +=1

Let's add currSum to the if clause as well, and rearrange things to look a little more like the else: clause.

    T = ['inf' for _ in range(len(nums))]
    count = 0
    for i in range(len(nums)):
        for j in range(i,len(nums)):
            if j == i:
                currSum = nums[j]
                T[j] = currSum 
                if currSum == k:
                    count +=1
            else:
                currSum = T[j-1] + nums[j]
                T[j] = currSum
                if currSum == k:
                    count +=1

Now we can see we can pull the if currSum == k: test out of the if...else clause.

    T = ['inf' for _ in range(len(nums))]
    count = 0
    for i in range(len(nums)):
        for j in range(i,len(nums)):
            if j == i:
                currSum = nums[j]
                T[j] = currSum 
            else:
                currSum = T[j-1] + nums[j]
                T[j] = currSum
            if currSum == k:
                count += 1

On the second and subsequent iterations of the inner loop, we add a number to T[j-1]. On the previous iteration, you stored currSum into that location. So, we can replace T[j-1] with currSum from the previous iteration.

    T = ['inf' for _ in range(len(nums))]
    count = 0
    for i in range(len(nums)):
        for j in range(i,len(nums)):
            if j == i:
                currSum = nums[j]
                T[j] = currSum 
            else:
                currSum = currSum + nums[j]
                T[j] = currSum
            if currSum == k:
                count += 1

At this point, you are storing into T[j] but never loading values from T[], so the entire T array can be deleted.

    count = 0
    for i in range(len(nums)):
        for j in range(i,len(nums)):
            if j == i:
                currSum = nums[j]
            else:
                currSum = currSum + nums[j]
            if currSum == k:
                count += 1

Finally, we can remove the first if statement by realizing it is just priming the accumulator currSum with 0 + nums[j] on the first iteration:

    count = 0
    for i in range(len(nums)):
        currSum = 0
        for j in range(i, len(nums)):
            currSum += nums[j]
            if currSum == k:
                count += 1

At this point, you can see your "dynamic programming" attempt was just busy work, and didn't actually improve your algorithm at all. In fact, it actually slowed it down and increased memory usage for no gain.

I have not changed your algorithm, just improved your implementation of the algorithm by removing the busy work. It is was and still is sequentially adding the numbers in range defined by a double for loop, so it still \$O(n^2)\$.

There is an improved algorithm which would use something like your T[] array. But this is a code-review, not an algorithm review, so I shan't spoil your fun of trying to find it.