C++で積がK未満となる部分配列の個数を数える方法

問題の概要

正の整数からなる配列 nums が与えられます。この中から、部分配列内の全要素の積が k 未満となる(連続した)部分配列の個数を数えて出力します。

例えば、入力が [10,5,2,6] で k = 100 の場合、出力は 8 になります。条件を満たす部分配列は以下の8つです。

[[10], [5], [2], [6], [10, 5], [5, 2], [2, 6], [5, 2, 6]]

解法のアプローチ

この問題はスライディングウィンドウ(二ポインタ)の手法を使うことで、O(n) の計算量で効率的に解くことができます。右端を固定しながら配列を走査し、積が k 以上になった時点で左端を縮めていくのがポイントです。

具体的な手順は以下の通りです。

• temp := 1、j := 0、ans := 0 で初期化します
• i を 0 から配列のサイズまで繰り返します
  • temp := temp * nums[i] とします
  • temp >= k かつ j <= i の間、以下を繰り返します
    • temp := temp / nums[j] とします
    • j を 1 増やします
  • ans := ans + (i - j + 1) とします
• ans を返します

ここで ans に加算する (i - j + 1) は、「右端が i である条件を満たす部分配列の個数」を表しています。ウィンドウ [j, i] 内のすべての部分配列が積 k 未満を保証されているためです。

C++の実装例

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

#include <bits/stdc++.h>
using namespace std;
typedef long long int lli;
class Solution {
public:
    int numSubarrayProductLessThanK(vector<int>& nums, int k) {
        lli temp = 1;
        int j = 0;
        int ans = 0;
        for(int i = 0; i < nums.size(); i++){
            temp *= nums[i];
            while(temp >= k && j <= i) {
                temp /= nums[j];
                j++;
            }
            ans += (i - j + 1);
        }
        return ans;
    }
};
main(){
    Solution ob;
    vector<int> v = {10,5,2,6};
    cout << (ob.numSubarrayProductLessThanK(v, 100));
}

入力

[10,5,2,6]
100

出力

8

計算量の評価

時間計算量は O(n) です。ポインタ i と j はそれぞれ最大で配列の長さ分しか進まないため、全要素を高々2回ずつ処理するだけで済みます。空間計算量は O(1) で、追加のデータ構造は不要です。また、積のオーバーフローを防ぐために temp を long long 型としている点にも注目してください。要素がすべて正の整数であるため、積は単調に増加し、この手法が正しく機能します。