numpyのクイックスタートチュートリアルにLess basic(あまり基本的でない)テクニックがあり面白そうと思ったので、理解ついでに紹介。
配列を用いた、配列のインデックス
配列のインデックス(a[i]におけるiの部分)にはスカラー値を持つことが一般的ですが、ここに配列を入れることもできます。
import numpy as np
a = np.arange(12)**2
i = np.array([1, 1, 3, 8, 5])
a[i] #array([ 1, 1, 9, 64, 25], dtype=int32)
何が起こっているのかというと、図のようになります。
配列iの要素がインデックスとなり、それを用いてaの配列から抽出してくるイメージでしょうか。
インデックスが2次元配列でも適用することができます。その場合、出力も2次元になります。
j = np.array([[3, 6, 7], [5, 9, 7]])
a[j]
#array([[ 9, 36, 49],
# [25, 81, 49]], dtype=int32)
チュートリアルの方ではRGBを応用例として出していますが、機械学習で使われるone-hot表現の際にも使えそうですね。
one_hot = np.array([[0, 0, 0],
[1, 0, 0],
[0, 1, 0],
[0, 0, 1]])
number = np.array([[0, 1, 2, 0],
[0, 3, 2, 0]])
one_hot[number]
#array([[[0, 0, 0],
# [1, 0, 0],
# [0, 1, 0],
# [0, 0, 0]],
#
# [[0, 0, 0],
# [0, 0, 1],
# [0, 1, 0],
# [0, 0, 0]]])
ちなみに、number[one_hot[number]]としても元に戻るわけではないので注意。
また、インデックスには複数の配列を指定することもできます。
a = np.arange(12).reshape(3,4)
#array([[ 0, 1, 2, 3],
# [ 4, 5, 6, 7],
# [ 8, 9, 10, 11]])
i = np.array([[0, 1],
[1, 2]])
j = np.array([[2, 1],
[3, 3]])
a[i, j]
#array([[ 2, 5],
# [ 7, 11]])
これまたどう処理しているのか解釈が難しいところですが、以下のようになっています。
配列のインデックスには、リストを指定することもできます。
a = np.arange(3,8)
a
#array([3, 4, 5, 6, 7])
a[[1,3,4]] = 0
a
#array([3, 0, 5, 0, 0])
このときも、リストの各要素がaのインデックスとして処理されます。
リストを使って一気に割り当てる(代入する)こともできますが、リスト内に同じ数値がある場合、割当ては繰り返されて最後の値が代入されます。
a = np.arange(3,8)
a
#array([3, 4, 5, 6, 7])
a[[1,1,4]] = [1,2,3]
a
#array([3, 2, 5, 6, 3])
ブール配列を使用したインデックス付け
配列に論理演算子を与えることで、ブール配列を作ることができます。
ブール配列をインデックスとすることで、Falseとなる要素を取り除いた1次元配列を出力します(配列の形に注意)。
a = np.arange(-3,9).reshape(3,4)
a
#array([[-3, -2, -1, 0],
# [ 1, 2, 3, 4],
# [ 5, 6, 7, 8]])
b = a > 0
b
#array([[False, False, False, False],
# [ True, True, True, True],
# [ True, True, True, True]])
a[b]
#array([1, 2, 3, 4, 5, 6, 7, 8])
ブール配列をインデックスとする配列に割り当てることで、条件に合った要素へ一度に代入することができます。
a[a<0] = 0
a
#array([[0, 0, 0, 0],
# [1, 2, 3, 4],
# [5, 6, 7, 8]])
次元(軸)と同じブール配列を使うことで、より複雑な抽出もできます。
a = np.arange(12).reshape(3,4)
b1 = np.array([False,True,True])
b2 = np.array([True,False,True,False])
a[b1,:] #a[b1]でも可
#array([[ 4, 5, 6, 7],
# [ 8, 9, 10, 11]])
a[:,b2] #a[b2]でも可
#array([[ 0, 2],
# [ 4, 6],
# [ 8, 10]])
a[b1,b2]
#array([ 4, 10])
図としてみるとこんなかんじ。
これ、なんでa[b1,b2]が[[4,6],[8,10]]じゃなくて[4,10]なんでしょうね。ドキュメントにもa weird thing to do(奇妙なこと)と書いてあるので、そう覚えるしかないのでしょうか。
まとめ
以上、配列を用いた配列のインデックスと、ブール配列を使用したインデックス付けについて紹介しました。扱いの難しいテクニックですが、使いこなせばきっと役に立つはずです。
冒頭にリンクを貼ったチュートリアルには、他にもテクニックが載っている(うまく理解できなかったため割愛)ので、余裕のある人はぜひ読んでみてください。