# Find sum of all the subarrays of an array

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?

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]


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))

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.