3430. Maximum and Minimum Sums of at Most Size K Subarrays

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  // Similar to 2104. Sum of Subarray Ranges
  long long minMaxSubarraySum(const vector<int>& nums, int k) {
    const auto [prevGt, nextGt] = getPrevNext(nums, less<>());
    const auto [prevLt, nextLt] = getPrevNext(nums, greater<>());
    return subarraySum(nums, prevGt, nextGt, k) +
           subarraySum(nums, prevLt, nextLt, k);
  }

 private:
  // Returns the sum of all subarrays with a size <= k, The `prev` and `next`
  // arrays are used to store the indices of the nearest numbers that are
  // smaller or larger than the current number, respectively.
  long subarraySum(const vector<int>& nums, const vector<int>& prev,
                   const vector<int>& next, int k) {
    long res = 0;
    for (int i = 0; i < nums.size(); ++i) {
      const int l = min(i - prev[i], k);
      const int r = min(next[i] - i, k);
      const int extra = max(0, l + r - 1 - k);
      res += nums[i] * static_cast<long>(l * r - extra * (extra + 1) / 2);
    }
    return res;
  }

  // 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};
  }
};

class Solution {
  // Similar to 2104. Sum of Subarray Ranges
  public long minMaxSubarraySum(int[] nums, int k) {
    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();
    return subarraySum(nums, k, prevGt, nextGt) + subarraySum(nums, k, prevLt, nextLt);
  }

  // Returns the sum of all subarrays with a size <= k, The `prev` and `next`
  // arrays are used to store the indices of the nearest numbers that are
  // smaller or larger than the current number, respectively.
  private long subarraySum(int[] nums, int k, int[] prev, int[] next) {
    long res = 0;
    for (int i = 0; i < nums.length; ++i) {
      final int l = Math.min(i - prev[i], k);
      final int r = Math.min(next[i] - i, k);
      final int extra = Math.max(0, l + r - 1 - k);
      res += (long) nums[i] * (l * r - extra * (extra + 1) / 2);
    }
    return res;
  }

  // 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);
  }
}

class Solution:
  # Similar to 2104. Sum of Subarray Ranges
  def minMaxSubarraySum(self, nums: list[int], k: int) -> int:
    prevGt, nextGt = self._getPrevNext(nums, operator.lt)
    prevLt, nextLt = self._getPrevNext(nums, operator.gt)
    return (self._subarraySum(nums, prevGt, nextGt, k) +
            self._subarraySum(nums, prevLt, nextLt, k))

  def _subarraySum(
      self,
      nums: list[int],
      prev: list[int],
      next: list[int],
      k: int
  ) -> int:
    """
    Returns the sum of all subarrays with a size <= k, The `prev` and `next`
    arrays are used to store the indices of the nearest numbers that are
    smaller or larger than the current number, respectively.
    """
    res = 0
    for i, num in enumerate(nums):
      l = min(i - prev[i], k)
      r = min(next[i] - i, k)
      extra = max(0, l + r - 1 - k)
      res += num * (l * r - extra * (extra + 1) // 2)
    return res

  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):
        index = stack.pop()
        next[index] = i
      if stack:
        prev[i] = stack[-1]
      stack.append(i)
    return prev, next