given an array I should compute the following sum in linear time:

my most naive implementation is O(N^2):

sum_ = 0

for i in range(n):
    for j in range(n, i, -1):
        sum_  = max(arr[i:j]) * (j-i)

I have no idea what to do? I have tried many algorithms but they were at best O(N*log(N)), but I should solve it in linear time, also, I don't get the idea, is there a mathematical way of just looking at an array and telling the result of the above sum?

CodePudding user response：

Keep a stack of (indices of) non-increasing values. So before appending the new value, pop smaller ones. Whenever you pop one, add its contribution to the total.

def solution(arr):
    arr.append(float('inf'))
    I = [-1]
    total = 0
    for i in range(len(arr)):
        while arr[i] > arr[I[-1]]:
            j = I.pop()
            a = j - I[-1]
            b = i - j
            total  = (a b)*a*b//2 * arr[j]
        I.append(i)
    arr.pop()
    return total

Correctness test with random lists of 500 values (Try it online!):

import random

def reference(arr):
    n = len(arr)
    return sum(max(arr[L : R 1]) * (R - (L-1))
               for L in range(n)
               for R in range(L, n))

for _ in range(5):
    arr = random.choices(range(10000), k=500)
    expect = reference(arr)
    result = solution(arr)
    print(result == expect, result)

Sample output (results for five lists, True means it's correct):

True 207276773131
True 208127393653
True 208653950227
True 208073567605
True 206924015682