C++でk個の奇数を含む「良い部分配列」の個数を数える方法

整数型の配列 nums と整数 k が与えられます。部分配列に含まれる奇数の個数がちょうど k 個であるとき、その部分配列を「良い部分配列(nice subarray)」と呼びます。この記事では、条件を満たす良い部分配列の個数を求めるアルゴリズムを解説します。

例として、配列が [1,1,2,1,1]、k = 3 の場合を考えてみましょう。このとき出力は 2 になります。条件を満たす部分配列は [1,1,2,1] と [1,2,1,1] の2つだからです。

解法の考え方

まず、配列中に現れる奇数のインデックスをすべて記録した配列 odd を作成します。次に、連続する k 個の奇数からなるウィンドウを順にスライドさせながら、各ウィンドウについて「左端の奇数より左側にいくつの開始位置を選べるか」「右端の奇数より右側にいくつの終了位置を選べるか」を求め、その積を答えに加算していきます。これにより、すべての条件を満たす部分配列を効率よく数えられます。

手順

• 答えを格納する変数 ans := 0、配列サイズ n := size(nums) で初期化します。
• left := 0、right := 0、count := 0 とします。
• 配列 odd を定義し、nums に含まれるすべての奇数のインデックスを格納します。
• odd の長さが k 以上である場合、以下を実行します。
  • i = 0、j = k - 1 から開始し、j が odd の末尾に達するまで i と j を1ずつ増やしながら繰り返します。
    • left := i が 0 なら odd[i] + 1、そうでなければ odd[i] - odd[i-1]
    • right := j が末尾なら n - odd[j]、そうでなければ odd[j+1] - odd[j]
    • ans := ans + left * right
• ans を返します。

C++ 実装例

以下の実装を見ると、処理の流れがより理解しやすくなります。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   int numberOfSubarrays(vector<int>& nums, int k) {
      int ans = 0;
      int n = nums.size();
      int left = 0;
      int right = 0;
      int cnt = 0;
      vector <int> odd;
      for(int i = 0; i < n; i++){
         if(nums[i] % 2 == 1)odd.push_back(i);
      }
      if(odd.size()>=k){
         for(int i = 0, j = k-1; j < odd.size(); i++, j++){
            int left = i==0?odd[i]+1: odd[i] - odd[i-1];
            int right = j==odd.size()-1 ?n-odd[j] : odd[j+1] - odd[j];
            ans += left * right;
         }
      }
      return ans;
   }
};
main(){
   vector<int> v = {1,1,2,1,1};
   Solution ob;
   cout <<ob.numberOfSubarrays(v, 3);
}

入力

[1,1,2,1,1]
3

出力

2

計算量

このアルゴリズムは、配列を一度走査して奇数のインデックスを収集した後、ウィンドウをスライドさせながら答えを累積していくため、時間計算量は O(n) です。また、奇数のインデックスを保存するための補助配列が必要となるため、空間計算量も O(n) となります。