C++で部分配列の最小値の合計を求める方法【単調スタックでO(N)高速化】

整数配列 A が与えられたとき、A のすべての（連続する）部分配列 B に対する min(B) の合計を求める問題を考えます。答えは非常に大きな値になる可能性があるため、10⁹ + 7 で割った余りを返します。

たとえば、入力が [3,1,2,4] の場合を考えてみましょう。部分配列は [3]、[1]、[2]、[4]、[3,1]、[1,2]、[2,4]、[3,1,2]、[1,2,4]、[3,1,2,4] の 10 個存在し、それぞれの最小値は [3,1,2,4,1,1,2,1,1,1] となります。これらの合計は 17 であるため、出力は 17 になります。

解法のアプローチ：単調スタック

すべての部分配列を列挙する方法では計算量が O(N²) となり、大きな入力には不向きです。そこで「単調スタック（モノトニックスタック）」を活用します。考え方の核心は、各要素 A[i] が「最小値」として寄与する部分配列の個数を求め、その個数に A[i] を掛けて合計に加算するというものです。

以下の手順で解きます。

• m := 1 × 10⁹ + 7 とします。
• 2 つのヘルパーメソッドを定義します。add(a, b) は (a mod m + b mod m) mod m を、mul(a, b) は (a mod m × b mod m) mod m を返します。
• メインメソッドでは配列 A を受け取り、スタック st を定義し、n := 配列 A のサイズとします。
• サイズ n の配列 left をすべて -1 で初期化し、同様にサイズ n の配列 right をすべて n で初期化します。
• ans := 0 とします。
• i を 0 から n − 1 までループします。
  • スタックが空でなく、A[スタックの先頭] ≥ A[i] である間、スタックから要素を取り除きます。
  • スタックが空でなければ、left[i] := スタックの先頭とします。
  • i をスタックにプッシュします。
• スタックが空になるまで要素を取り除きます。
• i を n − 1 から 0 まで逆順にループします。
  • スタックが空でなく、A[スタックの先頭] ≥ A[i] である間、スタックから要素を取り除きます。
  • スタックが空でなければ、right[i] := スタックの先頭とします。
  • i をスタックにプッシュします。
• i を 0 から n − 1 までループします。
  • leftBound := i − (left[i] + 1)、rightBound := (right[i] − 1) − i とします。
  • contri := 1 + leftBound + rightBound + (leftBound × rightBound) とします。
  • ans := add(ans, mul(contri, A[i])) とします。
• ans を返します。

なぜこの方法で正しく計算できるのか

left[i] は「A[i] より小さい要素が左側で最後に現れる位置」、right[i] は「右側で A[i] 未満の要素が最初に現れる位置」を表します。つまり、A[i] を最小値として含む部分配列の左端は leftBound 通り、右端は rightBound 通り選べるため、A[i] が最小値になる部分配列は合計で (leftBound + 1) × (rightBound + 1) = 1 + leftBound + rightBound + leftBound × rightBound 個存在します。

また、重複する値が含まれる場合に同じ部分配列を二重に数えないよう、左方向の走査では「以上（≥）」でポップし、右方向の走査では「より大きい（>）」でポップする、という非対称な条件を使っている点がポイントです。

例（C++）

理解を深めるために、以下の実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
const lli MOD = 1e9 + 7;
class Solution {
public:
    lli add(lli a, lli b){
        return (a % MOD + b % MOD) % MOD;
    }
    lli mul(lli a, lli b){
        return (a % MOD * b % MOD) % MOD;
    }
    int sumSubarrayMins(vector<int>& A) {
        stack <int> st;
        int n = A.size();
        vector <int> left(n, -1);
        vector <int> right(n, n);
        int ans = 0;
        for(int i = 0; i < n; i++){
            while(!st.empty() && A[st.top()] >= A[i]){
            st.pop();
        }
        if(!st.empty())left[i] = st.top();
            st.push(i);
        }
        while(!st.empty())st.pop();
        for(int i = n - 1; i >= 0; i--){
            while(!st.empty() && A[st.top()] > A[i]){
                st.pop();
            }
            if(!st.empty())right[i] = st.top();
                st.push(i);
        }
        for(int i = 0; i < n; i++){
            int leftBound = i - (left[i] + 1);
            int rightBound = (right[i] - 1) - i;
            int contri = 1 + leftBound + rightBound + (leftBound * rightBound);
            ans = add(ans, mul(contri, A[i]));
        }
        return ans;
    }
};
main(){
    vector<int> v = {3,1,2,4};
    Solution ob;
    cout << (ob.sumSubarrayMins(v));
}

入力

[3,1,2,4]

出力

17

このアルゴリズムでは、各要素がスタックにプッシュ・ポップされるのは高々 1 回ずつであるため、時間計算量 O(N)、空間計算量 O(N) で処理でき、総当たり方式と比べて大幅に高速化できます。