C++で行ごとにソートされた行列の中央値を効率的に求める方法

この記事では、行ごとにソートされた2次元配列 mat[r][c] が与えられたとき、その行列全体の要素の中から中央値（メディアン）を求めるアルゴリズムについて解説します。

問題の概要

各行が昇順にソートされている行列が与えられます。この行列に含まれるすべての要素の中で、中央に位置する値（中央値）を見つけることが目的です。

入力例

mat = {
    {2, 4, 7},
    {5, 6, 8},
    {4, 8, 9}
}

出力例

6

説明

行列のすべての要素を1つの配列にまとめると、以下のようになります。

{2, 4, 4, 5, 6, 7, 8, 8, 9}
中央値は 6

要素数は9個なので、ソート後の配列の5番目（中央）にある「6」が中央値となります。

解法アプローチ

シンプルな解法

最も単純な方法は、行列の全要素を1つの配列に格納し、それをソートした上で中央の要素を取り出すことです。実装は簡単ですが、時間計算量は O(r×c log(r×c)) となり、大規模な行列では非効率です。

二分探索を使った効率的な解法

より効果的なアプローチは、「中央値より小さい要素がちょうど (r×c)/2 個存在する」という性質を利用する方法です。具体的には以下の手順で進めます。

1. 行列内の最小値（各行の先頭要素の最小値）と最大値（各行の末尾要素の最大値）を求めます。
2. 最小値と最大値の範囲に対して二分探索を行います。
3. 範囲の中央値 mid を取り、mid より小さい（または等しい）要素の個数を数えます。
4. 個数が (r×c+1)/2 未満であれば、探索範囲の下限を mid + 1 に更新します。そうでなければ上限を mid に更新します。
5. 探索範囲が収束した時点の値が中央値です。

各要素の個数カウントには、C++の標準ライブラリ関数 upper_bound() を使うと便利です。これは各行がソート済みであることを利用して、mid より大きい最初の要素の位置を二分探索で高速に見つけられます。

この手法により、時間計算量は O(r × log(c) × log(max − min)) まで削減できます。

実装例

#include<bits/stdc++.h>
using namespace std;
#define c 3
#define r 3

int findMedian(int mat[][c]) {
    int smallest = INT_MAX, largest = INT_MIN;
    // 行列全体の最小値と最大値を求める
    for (int i = 0; i < r; i++) {
        if (mat[i][0] < smallest)
            smallest = mat[i][0];
        if (mat[i][c-1] > largest)
            largest = mat[i][c-1];
    }
    // 二分探索で中央値を特定する
    while (smallest < largest) {
        int mid = smallest + (largest - smallest) / 2;
        int smallCount = 0;
        // mid 以下の要素数を数える
        for (int i = 0; i < r; ++i)
            smallCount += upper_bound(mat[i], mat[i] + c, mid) - mat[i];
        if (smallCount < ((r * c + 1) / 2))
            smallest = mid + 1;
        else
            largest = mid;
    }
    return smallest;
}

int main() {
    int mat[][c] = { {2, 5, 7}, {4, 6, 8}, {1, 8, 9} };
    cout << "The median of the matrix is " << findMedian(mat);
    return 0;
}

出力

The median of the matrix is 6

まとめ

行ごとにソートされた行列の中央値を求める問題では、全要素をソートする代わりに、二分探索と upper_bound() を組み合わせることで大幅に計算量を抑えられます。特に大きな行列を扱う場合に有効なテクニックなので、ぜひ覚えておきましょう。