Pythonで奇数をk個含む「ナイスな部分配列」の個数を数えるプログラム

問題の概要

配列 nums と整数 k が与えられます。部分配列の中にちょうど k 個の奇数が含まれているとき、その部分配列を「ナイスな部分配列（nice subarray）」と呼ぶことにします。このとき、ナイスな部分配列が全体で何個存在するかを求めるのが本記事のテーマです。

たとえば、入力が nums = [1,1,2,1,1]、k = 3 の場合、出力は 2 になります。これは、[1,1,2,1] と [1,2,1,1] の2つの部分配列が、それぞれちょうど3つの奇数を含んでいるためです。

解決のためのアプローチ

この問題は、配列内の奇数が出現するインデックスをあらかじめ記録しておき、連続する k 個の奇数ブロックごとに「選べる始点の数 × 選べる終点の数」を掛け合わせて集計することで、効率よく解くことができます。

アルゴリズムの手順

1. 新しい空リスト odd_i を作成します。
2. i を 0 から nums のサイズ − 1 まで順に確認し、nums[i] % 2 == 1（nums[i] が奇数）であれば、そのインデックス i を odd_i の末尾に追加します。
3. start := 0、end := k − 1 と初期化します。
4. i := 0、count := 0 と初期化します。
5. end が odd_i のサイズ未満である間、次の処理を繰り返します。
   • end が odd_i のサイズ − 1 と等しい場合は j := len(nums) − 1 とし、そうでなければ j := odd_i[end + 1] − 1 とします。
   • count := count + (odd_i[start] − i + 1) * (j − odd_i[end] + 1) を計算して加算します。
   • i := odd_i[start] + 1、start := start + 1、end := end + 1 と更新します。
6. 最終的な count を返します。

なぜこの式で正しく数えられるのか

(odd_i[start] − i + 1) は、注目している k 個の奇数ブロックについて「左側に取りうる始点の候補数」を、(j − odd_i[end] + 1) は「右側に取りうる終点の候補数」を表しています。両者を掛け合わせることで、そのブロックを必ず含み、かつ余分な奇数を含まないすべての部分配列を、漏れや重複なく数え上げることができます。

Pythonでの実装例

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

def solve(nums, k):
   odd_i = []
   for i in range(len(nums)):
      if nums[i] % 2 == 1:
         odd_i.append(i)
   start = 0
   end = k - 1
   i = 0
   count = 0
   while end < len(odd_i):
      if end == len(odd_i) - 1:
         j = len(nums) - 1
      else:
         j = odd_i[end + 1] - 1
      count = count + (odd_i[start] - i + 1) * (j - odd_i[end] + 1)
      i = odd_i[start] + 1
      start = start + 1
      end = end + 1
   return count

nums = [1,1,2,1,1]
k = 3
print(solve(nums, k))

入力

nums = [1,1,2,1,1]
k = 3

出力

2

計算量

時間計算量：O(n) ― 配列の前処理と集計で、それぞれ配列を高々一度ずつ走査するだけです。
空間計算量：O(n) ― 奇数のインデックスを保存するためのリストが必要になります。