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
32
33
34
35
36
37
class Solution {
 public:
  long long subArrayRanges(vector<int>& nums) {
    const auto [prevGt, nextGt] = getPrevNext(nums, less<>());
    const auto [prevLt, nextLt] = getPrevNext(nums, greater<>());
    long ans = 0;

    for (int i = 0; i < nums.size(); ++i) {
      ans += static_cast<long>(nums[i]) * (i - prevGt[i]) * (nextGt[i] - i);
      ans -= static_cast<long>(nums[i]) * (i - prevLt[i]) * (nextLt[i] - i);
    }

    return ans;
  }

 private:
  // Returns `prev` and `next`, that store the indices of the nearest numbers
  // that are smaller or larger than the current number depending on `op`.
  pair<vector<int>, vector<int>> getPrevNext(
      const vector<int>& nums, const function<bool(int, int)>& op) {
    const int n = nums.size();
    vector<int> prev(n, -1);
    vector<int> next(n, n);
    stack<int> stack;
    for (int i = 0; i < n; ++i) {
      while (!stack.empty() && op(nums[stack.top()], nums[i])) {
        const int index = stack.top();
        stack.pop();
        next[index] = i;
      }
      if (!stack.empty())
        prev[i] = stack.top();
      stack.push(i);
    }
    return {prev, next};
  }
};

 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) {
    Pair<int[], int[]> gt = getPrevNext(nums, (a, b) -> a < b);
    Pair<int[], int[]> lt = getPrevNext(nums, (a, b) -> a > b);
    int[] prevGt = gt.getKey();
    int[] nextGt = gt.getValue();
    int[] prevLt = lt.getKey();
    int[] nextLt = lt.getValue();
    long ans = 0;

    for (int i = 0; i < nums.length; ++i) {
      ans += (long) nums[i] * (i - prevGt[i]) * (nextGt[i] - i);
      ans -= (long) nums[i] * (i - prevLt[i]) * (nextLt[i] - i);
    }

    return ans;
  }

  // Returns `prev` and `next`, that store the indices of the nearest numbers
  // that are smaller or larger than the current number depending on `op`.
  private Pair<int[], int[]> getPrevNext(int[] nums, BiFunction<Integer, Integer, Boolean> op) {
    final int n = nums.length;
    int[] prev = new int[n];
    int[] next = new int[n];
    Arrays.fill(prev, -1);
    Arrays.fill(next, n);
    Deque<Integer> stack = new ArrayDeque<>();
    for (int i = 0; i < n; ++i) {
      while (!stack.isEmpty() && op.apply(nums[stack.peek()], nums[i]))
        next[stack.pop()] = i;
      if (!stack.isEmpty())
        prev[i] = stack.peek();
      stack.push(i);
    }
    return new Pair<>(prev, next);
  }
}

 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
class Solution:
  def subArrayRanges(self, nums: list[int]) -> int:
    prevGt, nextGt = self._getPrevNext(nums, operator.lt)
    prevLt, nextLt = self._getPrevNext(nums, operator.gt)
    return sum(num * (i - prevGt[i]) * (nextGt[i] - i) -
               num * (i - prevLt[i]) * (nextLt[i] - i)
               for i, num in enumerate(nums))

  def _getPrevNext(
      self,
      nums: list[int],
      op: callable
  ) -> tuple[list[int], list[int]]:
    """
    Returns `prev` and `next`, that store the indices of the nearest numbers
    that are smaller or larger than the current number depending on `op`.
    """
    n = len(nums)
    prev = [-1] * n
    next = [n] * n
    stack = []
    for i, num in enumerate(nums):
      while stack and op(nums[stack[-1]], num):
        next[stack.pop()] = i
      if stack:
        prev[i] = stack[-1]
      stack.append(i)
    return prev, next