Skip to content

2104. Sum of Subarray Ranges 👍

  • 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
class Solution {
 public:
  long long subArrayRanges(vector<int>& nums) {
    return subarraySum(nums, less<int>()) - subarraySum(nums, greater<>());
  }

 private:
  long subarraySum(const vector<int>& A, const function<int(int, int)>& op) {
    const int n = A.size();
    long res = 0;
    vector<int> prev(n, -1);
    vector<int> next(n, n);
    stack<int> stack;

    for (int i = 0; i < n; ++i) {
      while (!stack.empty() && op(A[stack.top()], A[i])) {
        const int index = stack.top();
        stack.pop();
        next[index] = i;
      }
      if (!stack.empty())
        prev[i] = stack.top();
      stack.push(i);
    }

    for (int i = 0; i < n; ++i)
      res += static_cast<long>(A[i]) * (i - prev[i]) * (next[i] - i);

    return res;
  }
};

 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
class Solution {
  public long subArrayRanges(int[] nums) {
    return subarraySum(nums, '<') - subarraySum(nums, '>');
  }

  private long subarraySum(int[] A, char op) {
    final int n = A.length;
    long res = 0;
    int[] prev = new int[n];
    int[] next = new int[n];
    Deque<Integer> stack = new ArrayDeque<>();

    Arrays.fill(prev, -1);
    Arrays.fill(next, n);

    for (int i = 0; i < n; ++i) {
      while (!stack.isEmpty() && func(op, A[stack.peek()], A[i])) {
        final int index = stack.pop();
        next[index] = i;
      }
      if (!stack.isEmpty())
        prev[i] = stack.peek();
      stack.push(i);
    }

    for (int i = 0; i < n; ++i)
      res += (long) A[i] * (i - prev[i]) * (next[i] - i);

    return res;
  }

  private boolean func(char op, int a, int b) {
    if (op == '<')
      return a < b;
    return a > b;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution:
  def subArrayRanges(self, nums: List[int]) -> int:
    n = len(nums)

    def sumSubarray(A: List[int], op):
      res = 0
      prev = [-1] * n
      next = [n] * n
      stack = []

      for i, a in enumerate(A):
        while stack and op(A[stack[-1]], a):
          index = stack.pop()
          next[index] = i
        if stack:
          prev[i] = stack[-1]
        stack.append(i)

      for i, a in enumerate(A):
        res += a * (i - prev[i]) * (next[i] - i)

      return res

    return sumSubarray(nums, operator.lt) - sumSubarray(nums, operator.gt)