JavaScriptで異なる要素をちょうどn個含む部分配列の数を求める方法

問題

リテラル値からなる配列 arr を第1引数に、整数 num を第2引数に受け取る JavaScript の関数を作成する必要があります。この関数は、元の配列の中から「異なる要素をちょうど num 個含む連続した部分配列(サブ配列)」の個数を数え、その結果を返します。

たとえば、関数への入力が次の場合を考えてみましょう。

const arr = [12, 15, 12, 15, 18];
const num = 2;

このとき、出力は次のようになります。

const output = 7;

出力の解説

「異なる要素がちょうど 2 個」となる部分配列は、以下の 7 つです。

[12,15], [15,12], [12,15], [15,18],
[12,15,12], [15,12,15], [12,15,12,15]

アプローチ:スライディングウィンドウ

「ちょうど num 個の異なる要素を含む部分配列の数」は、次の関係式を利用すると効率よく求められます。

ちょうど num 個 = 「最大でも num 種類」の部分配列の数 −「最大でも (num − 1) 種類」の部分配列の数

「最大でも k 種類の異なる要素を含む部分配列の数」は、ハッシュマップで各要素の出現回数を記録しながら、右端を伸ばし・条件を満たさなくなったら左端を縮めるスライディングウィンドウ(二重ポインタ)の手法で O(n) 時間で計算できます。各ステップで有効なウィンドウの数は right - ptr + 1 として加算していきます。

コード例

const arr = [12, 15, 12, 15, 18];
const num = 2;
const distinctSubarrays = (arr = [], num = 1) => {
    const findDistinct = (count) => {
        const map = {};
        let ptr = 0;
        let distinct = 0;
        let res = 0;
        for(let right = 0; right < arr.length; right++){
            const num = arr[right];
            map[num] = (map[num] || 0) + 1;
            if(map[num] === 1){
                distinct += 1;
            };
            while(distinct > count){
                map[arr[ptr]] -= 1;
                if(map[arr[ptr]] === 0){
                    distinct -= 1;
                };
                ptr += 1;
            };
            res += right - ptr + 1;
        };
        return res;
    };
    return findDistinct(num) - findDistinct(num - 1)
};
console.log(distinctSubarrays(arr, num));

内部関数 findDistinct(count) は「異なる要素が最大でも count 種類である部分配列の総数」を返します。全体としては findDistinct(num) - findDistinct(num - 1) を計算することで、目的である「ちょうど num 種類」の部分配列の個数を導き出しています。

出力

コンソールには次のように表示されます。

7