Pythonで行列内の最大値を含むセルの個数を求めるプログラム

問題の概要

すべての要素が 0 で初期化された n × m の行列があるとします。ここに、「行と列の位置」のペアを格納したリストが与えられます。リスト内の各要素 i について、行番号と列番号がその要素の行の値・列の値より小さいすべてのセルの値が 1 ずつ加算されます。すべてのリスト要素を処理し終えた後、行列内で最大値を含むセルの個数を求める必要があります（行・列のインデックスは 0 から始まります）。

たとえば、入力が input_list = [[3, 5], [4, 6], [5, 3]] の場合、出力は 9 になります。ここでは 5 × 6 の行列を例に、処理の流れを確認してみましょう。

処理のステップを可視化

初期状態では、行列のすべての値は 0 です。

0 0 0 0 0 0
0 0 0 0 0 0
0 0 0 0 0 0
0 0 0 0 0 0
0 0 0 0 0 0

リストの最初の要素 [3, 5] を処理すると、行列は次のように変化します。

1 1 1 1 1 0
1 1 1 1 1 0
1 1 1 1 1 0
0 0 0 0 0 0
0 0 0 0 0 0

続いて 2 番目の要素 [4, 6] を処理すると、次のようになります。

2 2 2 2 2 1
2 2 2 2 2 1
2 2 2 2 2 1
1 1 1 1 1 1
0 0 0 0 0 0

最後に 3 番目の要素 [5, 3] を処理すると、行列は次の状態になります。

3 3 3 2 2 1
3 3 3 2 2 1
3 3 3 2 2 1
2 2 2 1 1 1
1 1 1 0 0 0

この時点で行列の最大値は 3 であり、その値を持つセルは 9 個あることが分かります。

解法の考え方

各操作のたびに行列全体を実際に更新し、最後に最大値を数える方法も思いつきますが、もっと賢いアプローチがあります。

鍵となるのは、セル (i, j) の最終的な値が「i 未満の行 r と j 未満の列 c を持つペアの個数」に等しいという点です。つまり、最大値を持つセルは、すべてのペアの条件を同時に満たす領域――すなわち「行の最小値 × 列の最小値」で囲まれた長方形の範囲――に必ず存在します。したがって、答えは min(r) × min(c) として一発で計算できます。

アルゴリズムの手順

• xpos := 0、ypos := 0 で初期化する
• input_list の各要素について、次を実行する
  • xpos が 0 の場合：xpos := item[0]、ypos := item[1] とする
  • それ以外の場合：xpos := min(xpos, item[0])、ypos := min(ypos, item[1]) とする
• xpos × ypos を返す

Pythonでの実装例

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

def solve(input_list):
    xpos = 0
    ypos = 0
    for item in input_list:
        if xpos == 0:
            xpos = item[0]
            ypos = item[1]
        else:
            xpos = min(xpos, item[0])
            ypos = min(ypos, item[1])
    return (xpos * ypos)

print(solve([[3, 5], [4, 6], [5, 3]]))

入力

[[3, 5], [4, 6], [5, 3]]

出力

9

計算量の評価

この解法はリストを一度だけ走査すればよいため、時間計算量は O(k)（k はペアの数）、追加メモリは O(1) で済みます。一方、行列を実際に構築して毎回更新する素朴な方法では O(行数 × 列数 × k) のコストがかかるため、行列のサイズが大きくなっても、この手法なら高速に動作します。