Skip to content

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<int>& strength) {
    constexpr int kMod = 1'000'000'007;
    const int n = strength.size();
    vector<long> prefix(n);
    vector<long> prefixOfPrefix(n + 1);

    for (int i = 0; i < n; ++i)
      prefix[i] = i == 0 ? strength[0] : (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<int> left(n, -1);
    stack<int> 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<int> right(n, n);
    stack = std::stack<int>();

    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 = 1_000_000_007;
    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[0] : (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<Integer> 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
36
class Solution:
  def totalStrength(self, strength: List[int]) -> int:
    kMod = 1_000_000_007
    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(itertools.accumulate(
        itertools.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