bluecode-小B的极差之和

问题描述

小B拥有一个数组 a，她使用这个数组构造了一个新数组 b。其中，a[i] 表示在新数组 b 中有 a[i] 个 i+1。例如，若 a = [2, 3, 1]，那么新数组 b = [1, 1, 2, 2, 2, 3]，因为 a[0] = 2 代表数组 b 中有 2 个 1，a[1] = 3 代表数组 b 中有 3 个 2，a[2] = 1 代表数组 b 中有 1 个 3。

现在，你需要帮助小B求出 b 数组中所有连续子数组的极差之和。由于答案可能非常大，请对 109+7109+7 取模。

数组的极差定义为子数组的最大值减去最小值。

测试样例

样例1：

输入：n = 2,a = [2, 1]
输出：2

样例2：

输入：n = 3,a = [1, 2, 1]
输出：6

样例3：

输入：n = 4,a = [2, 3, 1, 1]
输出：26

#include <iostream>
#include <stack>
#include <vector>

using namespace std;

const int MOD = 1e9 + 7;

long long get_contribute(const vector<int> &b, bool is_max) {
  int n = b.size();
  vector<int> left(n), right(n);
  stack<int> stk;

  // 找左边界
  for (int i = 0; i < n; ++i) {
    while (!stk.empty() &&
           (is_max ? b[stk.top()] <= b[i] : b[stk.top()] >= b[i]))
      stk.pop();
    left[i] = stk.empty() ? -1 : stk.top();
    stk.push(i);
  }

  // 清空栈
  while (!stk.empty())
    stk.pop();

  // 找右边界
  for (int i = n - 1; i >= 0; --i) {
    while (!stk.empty() && (is_max ? b[stk.top()] < b[i] : b[stk.top()] > b[i]))
      stk.pop();
    right[i] = stk.empty() ? n : stk.top();
    stk.push(i);
  }

  long long res = 0;
  for (int i = 0; i < n; ++i) {
    long long cnt = 1LL * (i - left[i]) * (right[i] - i) % MOD;
    res = (res + cnt * b[i]) % MOD;
  }

  return res;
}

int solution(int n, vector<int> a) {
  vector<int> b;
  for (int i = 0; i < n; ++i) {
    b.insert(b.end(), a[i], i + 1);
  }

  long long maxSum = get_contribute(b, true);
  long long minSum = get_contribute(b, false);

  return (maxSum - minSum + MOD) % MOD;
}

int main() {
  cout << (solution(2, {2, 1}) == 2) << endl;
  cout << (solution(3, {1, 2, 1}) == 6) << endl;
  cout << (solution(4, {2, 3, 1, 1}) == 26) << endl;
  return 0;
}