1
\$\begingroup\$

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?

\$\endgroup\$
0
4
\$\begingroup\$

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.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.