Pythonで解く：ソート済み行列におけるK番目に小さい要素の求め方

問題概要

各行・各列が昇順にソートされた n × n の行列が与えられたとき、行列内で k 番目に小さい要素を求めることを考えます。ここで注意すべきは、求めるのは「ソート順における k 番目の要素」であり、「k 番目に小さい一意な値」ではないという点です。つまり、同じ値が複数回現れる場合は、それぞれを別々にカウントします。

例えば、入力が [[1,5,9],[10,11,13],[12,13,15]] で k = 8 の場合、出力は 13 になります。

アプローチ：二分探索との組み合わせ

この問題は、要素そのものを並べ替えるのではなく、「ある値以下の要素がいくつあるか」というカウント結果に対して二分探索を行うことで、効率的に解くことができます。

手順1：checkVal() メソッドの定義

まず、指定した値以下の要素数を数える checkVal() メソッドを定義します。引数は matrix と value です。

1. i := 0、j := matrix[0] の長さ − 1、counter := 0 と初期化します。
2. i < 行列の行数 かつ j ≥ 0 の間、以下を繰り返します。
   • matrix[i][j] > value ならば、j を 1 減らします。
   • そうでなければ、counter に j + 1 を加算し、i を 1 増やします。
3. counter を返します。これが「value 以下の要素の総数」です。

この走査は行列の右上隅から開始し、階段状に移動しながら条件を満たす要素を数えるため、わずか O(n) の時間で完了します。

手順2：メインメソッドでの二分探索

• n := 行列の行数、high := 右下隅の要素（最大値）、low := 左上隅の要素（最小値）とします。
• low ≤ high の間、以下を繰り返します。
  • mid = low + (high − low) / 2
  • count := checkVal(matrix, mid)
  • count < k ならば low := mid + 1、そうでなければ high := mid − 1
• low を返します。これが k 番目に小さい要素です。

計算量

checkVal() が O(n)、二分探索が O(log(max − min)) 回繰り返されるため、全体の時間計算量は O(n log(max − min))、空間計算量は O(1) となります。行列をフラット化してソートする O(n² log n) の素朴な手法と比べ、大幅に高速です。

実装例

class Solution(object):
    def kthSmallest(self, matrix, k):
        """
        :type matrix: List[List[int]]
        :type k: int
        :rtype: int
        """
        n = len(matrix)
        high = matrix[n-1][n-1]
        low = matrix[0][0]
        while low <= high:
            mid = low + (high - low) / 2
            count = self.check_value(matrix, mid)
            if count < k:
                low = mid + 1
            else:
                high = mid - 1
        return int(low)

    def check_value(self, matrix, value):
        i = 0
        j = len(matrix[0]) - 1
        counter = 0
        while(i < len(matrix) and j >= 0):
            if matrix[i][j] > value:
                j -= 1
            else:
                counter += j + 1
                i += 1
        return counter

matrix = [[1,5,9],[10,11,13],[12,13,15]]
ob = Solution()
print(ob.kthSmallest(matrix, 8))

入力

matrix = [[1,5,9],[10,11,13],[12,13,15]]
k = 8

出力

13

まとめ

ソート済み行列から k 番目に小さい要素を求める問題は、「値以下の個数を数える関数」と「答えの候補値に対する二分探索」を組み合わせることで、O(n log(max − min)) という高い効率で解決できます。右上隅から階段状に走査するテクニックは、ソート済み行列を扱う多くの問題に応用できるので、ぜひ覚えておきましょう。