# 2281. Sum of Total Strength of Wizards      • Time: $O(n)$
• Space: $O(n)$
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 class Solution { public: int totalStrength(vector& strength) { constexpr int kMod = 1e9 + 7; const int n = strength.size(); vector prefix(n); vector prefixOfPrefix(n + 1); for (int i = 0; i < n; ++i) prefix[i] = i == 0 ? strength : (strength[i] + prefix[i - 1]) % kMod; for (int i = 0; i < n; ++i) prefixOfPrefix[i + 1] = (prefixOfPrefix[i] + prefix[i]) % kMod; // next small or equal on the left vector left(n, -1); stack stack; for (int i = n - 1; i >= 0; --i) { while (!stack.empty() && strength[stack.top()] >= strength[i]) left[stack.top()] = i, stack.pop(); stack.push(i); } // next small on the right vector right(n, n); stack = std::stack(); for (int i = 0; i < n; ++i) { while (!stack.empty() && strength[stack.top()] > strength[i]) right[stack.top()] = i, stack.pop(); stack.push(i); } long ans = 0; // for each strength[i] as minimum, calculate sum for (int i = 0; i < n; ++i) { const int l = left[i]; const int r = right[i]; const long leftSum = prefixOfPrefix[i] - prefixOfPrefix[max(0, l)]; const long rightSum = prefixOfPrefix[r] - prefixOfPrefix[i]; const int leftLen = i - l; const int rightLen = r - i; ans += strength[i] * (rightSum * leftLen % kMod - leftSum * rightLen % kMod + kMod) % kMod; ans %= kMod; } return ans; } }; 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 class Solution { public int totalStrength(int[] strength) { final int kMod = (int) 1e9 + 7; final int n = strength.length; long[] prefix = new long[n]; long[] prefixOfPrefix = new long[n + 1]; for (int i = 0; i < n; ++i) prefix[i] = i == 0 ? strength : (strength[i] + prefix[i - 1]) % kMod; for (int i = 0; i < n; ++i) prefixOfPrefix[i + 1] = (prefixOfPrefix[i] + prefix[i]) % kMod; // next small or equal on the left int[] left = new int[n]; Arrays.fill(left, -1); Deque stack = new ArrayDeque<>(); for (int i = n - 1; i >= 0; --i) { while (!stack.isEmpty() && strength[stack.peek()] >= strength[i]) left[stack.pop()] = i; stack.push(i); } // next small on the right int[] right = new int[n]; Arrays.fill(right, n); stack = new ArrayDeque<>(); for (int i = 0; i < n; ++i) { while (!stack.isEmpty() && strength[stack.peek()] > strength[i]) right[stack.pop()] = i; stack.push(i); } long ans = 0; // for each strength[i] as minimum, calculate sum for (int i = 0; i < n; ++i) { final int l = left[i]; final int r = right[i]; final long leftSum = prefixOfPrefix[i] - prefixOfPrefix[Math.max(0, l)]; final long rightSum = prefixOfPrefix[r] - prefixOfPrefix[i]; final int leftLen = i - l; final int rightLen = r - i; ans += strength[i] * (rightSum * leftLen % kMod - leftSum * rightLen % kMod + kMod) % kMod; ans %= kMod; } return (int) ans; } } 
  1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 class Solution: def totalStrength(self, strength: List[int]) -> int: kMod = int(1e9) + 7 n = len(strength) # next small or equal on the left left = [-1] * n stack = [] for i in reversed(range(n)): while stack and strength[stack[-1]] >= strength[i]: left[stack.pop()] = i stack.append(i) # next small on the right right = [n] * n stack = [] for i in range(n): while stack and strength[stack[-1]] > strength[i]: right[stack.pop()] = i stack.append(i) ans = 0 prefixOfPrefix = list(accumulate(accumulate(strength), initial=0)) # for each strength[i] as minimum, calculate sum for i, (l, r) in enumerate(zip(left, right)): leftSum = prefixOfPrefix[i] - prefixOfPrefix[max(0, l)] rightSum = prefixOfPrefix[r] - prefixOfPrefix[i] leftLen = i - l rightLen = r - i ans += strength[i] * (rightSum * leftLen - leftSum * rightLen) % kMod return ans % kMod