JavaScriptでソート済み2次元配列からN番目に小さい要素を求める方法

問題の概要

昇順にソートされた数値の2次元配列(行列)が与えられたとします。

const arr = [
    [ 1, 5, 9],
    [10, 11, 13],
    [12, 13, 15]
];

この配列を第1引数として受け取り、整数 num を第2引数として受け取るJavaScript関数を作成します。関数は、配列 arr 内に存在する num 番目に小さい要素を返す必要があります。

例えば、関数への入力が以下の場合を考えてみましょう。

const arr = [
    [ 1, 5, 9],
    [10, 11, 13],
    [12, 13, 15]
];
const num = 5;

出力

const output = 11;

出力の説明

行列の要素を小さい順に並べると「1, 5, 9, 10, 11, ...」となるため、11は5番目に小さい要素です。

実装コード

この処理を実装したコードは以下の通りです。

const arr = [
    [ 1, 5, 9],
    [10, 11, 13],
    [12, 13, 15]
];
const num = 5;

const kthSmallest = (arr = [], num = 1) => {
    let low = arr[0][0];
    let high = arr[arr.length - 1][arr[0].length - 1] + 1;
    while (low < high) {
        let mid = low + Math.floor((high - low) / 2);
        let count = 0;
        for (let i = 0; i < arr.length; i++) {
            for (let j = 0; j < arr.length; j++) {
                if (arr[i][j] <= mid) count++;
                else break;
            }
        }
        if (count < num) low = mid + 1;
        else high = mid;
    }
    return low;
};

console.log(kthSmallest(arr, num));

コード解説

このアルゴリズムの基本的な考え方を順を追って説明します。

1次元配列との違い

通常のソート済み1次元配列であれば、インデックスの範囲を使って目的の値を二分探索できます。例えば以下のようなイメージです。

low = 0, high = length - 1, mid = (low + high) / 2

目的の値が mid 番目の要素より大きければ右側を、小さければ左側を探索します。

しかし、今回のような2次元配列では、そのような「mid」となる単一のインデックスを見つけることができません。そこで発想を転換し、「値そのものの範囲」を使ってk番目の要素を絞り込んでいくのがポイントです。

値の範囲による二分探索

配列の最初の要素が最小値、最後の要素が最大値であることは分かっているため、答えとなる数値は必ずその間に存在します。この2つの値を low と high として設定し、mid をその中間の値とします。そして、配列内に mid 以下の値がいくつあるかをカウントし、その結果に応じて low と high を調整していきます。

• count が num 未満なら、答えはもっと大きい値 → low = mid + 1
• count が num 以上なら、答えは mid 以下 → high = mid

low と high が収束したとき、それこそが求めていたk番目に小さい要素です。

midが配列内に存在しない場合の扱い

ここで注意したいのは、mid の計算方法からすると、テスト対象の数値が必ずしも配列内に存在するとは限らないという点です。結局のところ、使用しているのは任意の数値だからです。

これを理解するには、プログラムが low と high がほぼ衝突する段階に達した場面を想像してください。期待より少ない数しかカウントできなかった場合は low = mid + 1 と設定されるため、mid が1ずつ増加していきます。このように1ずつ値を詰めていくことで、最終的には必ず配列内に実際に存在する数値に行き着く仕組みになっています。

実行結果

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

11