Find sum of all the subarrays of an array

Question

I have an array. Now I want to find the total of, sum of the subarray multiplied by its last element, for all the possible subarrays. This is what I have right now:

n = int(input())
a = []
for _ in range(n):
 a.append(int(input()))
total = 0
for i in range(1, n+1):
 for j in range(n+1-i):
 temp = a[j:j+i]
 total += sum(temp)*temp[-1]
print(total)

Example Input:

3
1
2
3

Output:

53

Explanation:

1*1 + 2*2 + 3*3 + (1+2)*2 + (2+3)*3 + (1+2+3)*3 = 53

My code works fine, but is quite slow. How can I optimise it?

Answer 1

First I would suggest to separate I/O from the computation, and define a function to compute the subarray sum. That increases the clarity of the program and allows to add test cases more easily:

def subarray_sum(a):
 """Compute sum of all subarrays of a, multiplied by its last element"""
 n = len(a)
 total = 0
 for i in range(1, n + 1):
 for j in range(n + 1 - i):
 temp = a[j:j + i]
 total += sum(temp) * temp[-1]
 return total

Using sum() with generator expressions this can be shortened to

def subarray_sum(a):
 n = len(a)
 total = sum(sum(sum(a[i:j + 1]) * a[j] for j in range(i, n))
 for i in range(n))
 return total

But the time complexity is still \$ O(n^3) \$ because of the three nested loops.

In order to find a more efficient method, let's compute the sum for a 3-element array \$ [a, b, c] \$ explicitly:

$$ a \cdot a + b \cdot b + c \cdot c \\ + (a+b)\cdot b + (b+c) \cdot c \\ + (a+b+c) \cdot c $$ Rearranging terms, this becomes $$ a \cdot a + (a + 2b) \cdot b + (a + 2b + 3c) \cdot c $$ Can you spot the pattern? This can be computed with a single traversal of the array, i.e. in \$ O(n) \$ time.