C++で有効な部分配列の個数を求めるアルゴリズム

整数の配列 A が与えられたとき、「部分配列の左端の要素が、その部分配列内の他のすべての要素以下である」という条件を満たす、空でない連続した部分配列の個数を求めます。

例えば、入力が [1,4,2,5,3] の場合、出力は 11 になります。条件を満たす部分配列は次の 11 個です。

[1], [4], [2], [5], [3], [1,4], [2,5], [1,4,2], [2,5,3], [1,4,2,5], [1,4,2,5,3]

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

この問題は「単調スタック(増加スタック)」を使うことで効率的に解けます。各要素を処理する際、スタックの先頭にある自分より大きい要素をすべて取り除いてから、現在の要素をプッシュします。このときのスタックのサイズが、「現在の位置で終わる有効な部分配列の個数」に対応します。

手順

• 答えを格納する変数 ret を 0 で初期化します。
• n を配列 nums のサイズとします。
• スタック st を用意します。
• i を 0 から nums のサイズまで繰り返します。
  • x := nums[i] とします。
  • スタックが空でなく、かつ x がスタックの先頭要素より小さい間、スタックから要素を取り除きます(pop)。
  • x をスタックに挿入します(push)。
  • ret にスタックのサイズを加算します。
• 最後に ret を返します。

なぜこれで正しく求まるのか

スタックには「現在の要素より大きくない値」だけが残ります。つまり、スタック内の各要素は、そこを始点とし現在の位置を終点とする部分配列の左端として条件を満たしています。したがって、各時点でのスタックのサイズを合計すれば、すべての有効な部分配列を重複なく数えることができます。

C++での実装例

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

#include <bits/stdc++.h>
using namespace std;
class Solution {
    public:
    int validSubarrays(vector<int>& nums) {
        int ret = 0;
        int n = nums.size();
        stack <int> st;
        for(int i = 0; i < nums.size(); i++){
            int x = nums[i];
            while(!st.empty() && x < st.top()) st.pop();
            st.push(x);
            ret += st.size();
        }
        return ret;
    }
};
main(){
    Solution ob;
    vector<int> v = {1,4,2,5,3};
    cout << (ob.validSubarrays(v));
}

入力

{1,4,2,5,3}

出力

11

計算量

• 時間計算量: O(n)。各要素は最大でも 1 回プッシュされ、1 回しかポップされないためです。
• 空間計算量: O(n)。最悪ケース(昇順に並んだ配列など)では、スタックにすべての要素が保持されます。