numpy.diag – 2次元配列から対角線の要素を取得または1次元配列を対角線に配置

numpy.diagは、2次元配列の場合は対角線の要素を抽出した配列、1次元配列の場合はその要素を対角線に配置し、他を0埋めした配列を生成する関数です。

実際のコードで確認していきましょう。

1. numpy.diagの使い方

まずはnumpy.diagの基本的な書き方を確認しましょう。

numpy.diag(v, k=0)

それではサンプルコードを見ながら使い方を確認していきましょう。

1.1. 2次元配列を渡す場合

numpy.diagに2次元配列を渡すと、その配列の対角線の要素を格納した新しい配列を生成します。

まず、以下のコードで2次元配列を生成します。

'''  2次元配列を渡すと対角線の要素を格納した配列を生成  '''
import numpy as np
#  2次元配列を生成
arr = np.arange(9).reshape(3, 3)
arr

array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

この配列を、numpy.diag()に渡すと、対角線の要素を抽出した配列が返ってきます。

#  対角線の要素を格納した配列を生成
np.diag(arr)

array([0, 4, 8])

この時、オプション引数「k=」で、どの対角線の要素を取得するかを指定することができます。

以下のコードのように正の値を指定すると、指定の値だけ上から右にずれた対角線になります。

#  インデックス1の対角線（上から右にずれる）
np.diag(arr, k=1)

array([1, 5])

負の値を指定すると、指定の値だけ下から左にずれた対角線になります。

#  インデックス-1の対角線（下から左にずれる）
np.diag(arr, k=-1)

array([3, 7])

1.2. 1次元配列を渡す場合

1次元配列を渡すと、その配列の要素を対角線に配置し、他の要素を0埋めした2次元配列を返します。元の1次元配列がshape(N)なら、生成される2次元配列はshape(N, N)になります。

以下のコードでご確認ください。

'''  1次元配列を渡すと、その要素を対角線し、他を0埋めの配列を生成'''
arr = np.arange(5)
print('arr:\n', arr)
diag = np.diag(arr)
print('diag:\n', diag)

arr:
 [0 1 2 3 4]
diag:
 [[0 0 0 0 0]
 [0 1 0 0 0]
 [0 0 2 0 0]
 [0 0 0 3 0]
 [0 0 0 0 4]]

1.3. 元の配列の対角線以外の要素を0埋め

次のようにnumpy.diag()を重ねがけして、元の配列の対角線以外の要素を0埋めにするというような使い方もよく見られますので覚えておくと良いでしょう。

'''  次のように重ねがけする場合が多いでしょう  '''
#  配列を生成
arr = np.arange(1, 10).reshape(3, 3)
print('arr:\n', arr)
#  numpy.diag()の重ねがけで対角線の要素以外を0埋め
diag = np.diag(np.diag(arr))
print('diag:\n', diag)

arr:
 [[1 2 3]
 [4 5 6]
 [7 8 9]]
diag:
 [[1 0 0]
 [0 5 0]
 [0 0 9]]

2. まとめ

以上がnumpy.diag()の使い方です。

データサイエンスにおいて対角線の操作は頻繁に行うので、覚えておきましょう。また、以下の関数も確認しておくと良いでしょう。

• numpy.eye: 対角線の要素が1で他は0埋めの配列を生成
• numpy.identity: 対角線の要素が1で他は0埋めの正方形の配列を生成
• numpy.diagonal: 既存の配列の対角線を要素とした配列を生成
• numpy.diagflat: 対角線が指定の要素、他が0埋めの配列を生成
• numpy.trace: 配列の対角線の要素の合計を取得