subarray_sum.py

'''
This problem is taken from Stanford's CS 9: Problem-Solving for
                                            the CS Technical Interview.
http://web.stanford.edu/class/cs9/sample_probs/SubarraySums.pdf

A subarray of an array is a consecutive sequence of zero or more values taken
out of that array. For example, the array [1, 3, 7] has seven subarrays:
[ ] [1] [3] [7] [1, 3] [3, 7] [1, 3, 7]

Notice that [1, 7] is not a subarray of [1, 3, 7], because even though the
values 1 and 7 appear in the array, they're not consecutive in the array.

Similarly, the array [7, 3] isn't a subarray of the original array,
because these values are in the wrong order.

The sum of an array is the sum of all the values in that array. Your task is
to write a function that takes as input an array and outputs the sum of all of
its subarrays.

For example, given [1, 3, 7], you'd output 36, because:

[ ] + [1] + [3] + [7] + [1, 3] + [3, 7] + [1, 3, 7]
= 0 + 1 + 3 + 7 + 4 + 10 + 11 = 36

My solution:
    subarray_sum_slow: the easy but slow way to do it. Just enumberate all the
                       subarrays and sum them up.
    subarray_sum: the fast way to do it. We can calculate the number of times
                  each element occurs and sum them up.
'''


def get_subarrays(arr):
    subarrays = []
    for i in range(len(arr)):
        for j in range(i, len(arr)):
            subarrays.append(arr[i:j + 1])
    return subarrays


def subarray_sum_slow(arr):
    subarrays = get_subarrays(arr)
    return sum([sum(arr) for arr in subarrays])


def subarray_sum(arr):
    total = 0
    for i, num in enumerate(arr):
        total += num * (i + 1) * (len(arr) - i)
    return total


assert subarray_sum_slow([1]) == subarray_sum([1])
assert subarray_sum_slow([1, 2]) == subarray_sum([1, 2])
assert subarray_sum_slow([1, 2, 3]) == subarray_sum([1, 2, 3])
assert subarray_sum_slow([1, 3, 7]) == subarray_sum([1, 3, 7])
assert subarray_sum_slow([1, 3, 7, 9]) == subarray_sum([1, 3, 7, 9])
assert subarray_sum_slow([1, 3, 7, 9, 11]) == subarray_sum([1, 3, 7, 9, 11])
assert subarray_sum_slow(
    [1, 3, 7, 9, 11, 13]) == subarray_sum([1, 3, 7, 9, 11, 13])