Задача с собеседования в Amazon

Дан массив целых чисел. Необходимо посчитать сумму минимумов каждого подмассива исходного массива. Ответ вывести по модулю 1e9

Пример:
[3, 1, 2, 4]
Ответ: 17

Наш чат алгоритмистов

Решение:
Для каждого элемента A[i] посчитать, в скольких подмассивах он будет минимальным. Найдём индекс предыдущего элемента строго меньше (prev) и следующий элемент меньше или равный (next) с помощью монотонного стека. Тогда вклад A[i] в сумму равен A[i] * (i - prev[i]) * (next[i] - i). Время O(n), память O(n).

def sumSubarrayMins(A: List[int]) -> int:
n = len(A)
if n == 0:
return 0

# prev: индекс предыдущего элемента, который строго меньше A[i]
prev = [-1] * n
stack = []
for i in range(n):
while stack and A[stack[-1]] >= A[i]:
stack.pop()
prev[i] = stack[-1] if stack else -1
stack.append(i)

# nxt: индекс следующего элемента, который меньше или равен A[i]
nxt = [n] * n
stack = []
for i in range(n - 1, -1, -1):
while stack and A[stack[-1]] > A[i]:
stack.pop()
nxt[i] = stack[-1] if stack else n
stack.append(i)

res = 0
for i in range(n):
left_count = i - prev[i]
right_count = nxt[i] - i
res = (res + A[i] * left_count * right_count) % MOD

return res

@algoses