Pythonで行列の対角線上の要素の合計を求めるプログラム

問題の概要

正方行列が与えられたとき、その対角線上の要素の合計を求めることを考えます。具体的には、主対角線（左上から右下）と副対角線（右上から左下）に含まれるすべての要素を足し合わせます。ただし、奇数次の行列では2つの対角線が交差する中央の要素が重複するため、その要素は1回だけカウントする必要があります。

たとえば、次のような4×4の行列が入力として与えられたとします。

この場合の出力は 59 になります。内訳を見ると、主対角線の要素は [10, 15, 12, 3] でその合計は 40、副対角線の要素は [6, 3, 8, 2] でその合計は 19 です。したがって、全体の合計は 40 + 19 = 59 となります。

アルゴリズムの手順

この問題は、以下の手順に従うことで効率的に解くことができます。

• m を行列の行数とする。

• m が 1 の場合は、matrix[0][0] を返す。

• count を 0 で初期化する。

• i を 0 から m - 1 までループさせる。

  • count に matrix[i][i] を加算する。

  • count に matrix[i][-1 - i] を加算する。

• m が奇数の場合：

  • ind を m / 2 の商（小数点以下切り捨て）とする。

  • count から matrix[ind][ind] を減算し、重複分を取り除く。

• count を返す。

Pythonでの実装例

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

def solve(matrix):
    m = len(matrix)
    if m == 1: return matrix[0][0]

    count = 0
    for i in range(m):
        count += matrix[i][i]
        count += matrix[i][-1 - i]

    if m % 2 == 1: count -= matrix[m // 2][m // 2]

    return count

matrix = [[10,5,9,6],[8,15,3,2],[3,8,12,3],[2,11,7,3]]
print(solve(matrix))

入力

[[10,5,9,6],[8,15,3,2],[3,8,12,3],[2,11,7,3]]

出力

59

コードのポイント解説

この実装の鍵となるのは負のインデックスの活用です。Pythonでは matrix[i][-1 - i] と書くことで、各行の末尾から数えて i 番目の要素、つまり副対角線上の要素に簡単にアクセスできます。

また、行列のサイズ m が奇数の場合、中央の要素 matrix[m // 2][m // 2] は主対角線と副対角線の両方に属するため、ループ内で2回加算されたことになります。そこで最後に1回だけ減算することで、重複カウントを防いでいます。

計算量は、行列のすべての行を一度ずつ処理するだけなので O(m) となり、非常に効率的なアルゴリズムです。