多次元配列から任意の要素を取得するnumpy.choose()の使い方

numpy.choose()は、インダイスの配列によって、多次元配列から要素を取得する際に使います。これを使わない場合と比べて、かなり自由に任意の値を取得することが可能になります。

解説が多少長くなっていますので、まずサンプルコードを先に見てからの方が理解が早いかもしれません。それでは見ていきましょう。

1. numpy.choose()の使い方

まずは基本書式を確認しましょう。

np.choose(a, choices, out=None, mode='raise')

ndarray.choose(choices, out=None, mode='raise')

Note

この関数は基本的に以下のコードと同じです。

np.choose(a,choices) == np.array([c[a[i]][i] for i in ndi.ndindex(a.shape)])
*ndi = numpy.lib.index_tricks

基本的にはこのように覚えておけば問題ありません。

しかし、modeの指定によって、配列a のインダイスの数値が、choices に渡すシーケンス内の配列の数を超えている場合の扱いが異なります。この点を含めると、厳密には次のような手順で処理されます。

numpy.choose(a, choices)には、第一引数a に要素が整数(int型)のインデックス配列を渡し、第二引数choices に配列のシーケンスを渡します。すると、a 及び choices のそれぞれの配列は、必要ならまず、Ba と Bchoices[i], i = 0,…,n-1 を呼び出すことでブロードキャストします。必然的に、それぞれのi についてBa.shape == Bchoices[i].shapeとなります。次に Ba.shape の形状の新しい配列が以下のルールに則って生成されます。

2. サンプルコード

それではサンプルコードを確認しましょう。

まず、第一引数に渡す配列a と第二引数に渡す配列choices を作成しておきます。

import numpy as np
a = np.array([2, 3, 1, 0])
choices = np.array([[0, 1, 2, 3], [10, 11, 12, 13], [20, 21, 22, 23], [30, 31, 32, 33]])

numpy.choose()は、この配列choices の中から、第一引数に渡したインダイス配列a に従って順番に要素を取得していきます。詳しくは、以下のコード及び説明をしっかりご確認ください。

# 出力結果の最初の要素は、
# choicesの3つ目の配列の最初の要素の20に
# 出力結果の2つ目の要素は、
# choicesの4つ目の配列の2つ目の要素31に
# 出力結果の3つ目の要素は、
# choicesの2つ目の配列の3つ目の要素12に
# 出力結果の4つ目の要素は、
# choicesの最初の配列の4つ目の要素3になる
np.choose(a, choices)

array([20, 31, 12,  3])

numpy.choose()はデフォルト(mode=’raise’)の場合、以下のようにリスト内包表記で要素を取り出す操作と出力結果は同じです。

# 以下のコードと同じ
np.array([choices[a[i]][i] for i in np.ndindex(a.shape)])

array([20, 31, 12,  3])

しかし、mode を ‘clip’ または ‘wrap’にすると処理が変わります。

# modeを'clip'にすると、
# 範囲外のインダイス 4 を 3 に変換
np.choose([2, 4, 1, 0], choices, mode='clip')

array([20, 31, 12,  3])

# modeを'wrap'にすると、
# 範囲外のインダイス 4 を [4 mod 配列数] に変換
# mod は剰余演算（余りを求める演算のこと)
# この場合は 4 mod 4 = 0となる。
np.choose([2, 4, 1, 0], choices, mode='wrap')

array([20,  1, 12,  3])

このように numpy.choose() は、choices の要素の中から、インダイス配列a の指定にしたがって順々に要素を取得して新しい配列を作成します。

以下はブロードキャストを使った場合のサンプルコードです。ブロードキャストについては『覚えておくべきNumPy配列のブロードキャストのルール』で詳しく解説していますので是非ご確認ください。

import numpy as np
a = np.array([[1, 0, 1], [0, 1, 0], [1, 0, 1]])
a

array([[1, 0, 1],
       [0, 1, 0],
       [1, 0, 1]])

choices = np.array([-10, 10])
choices

array([-10,  10])

# 1の場所に10を、0の場所に-10を取得
np.choose(a, choices)

array([[ 10, -10,  10],
       [-10,  10, -10],
       [ 10, -10,  10]])

もう1つ高度な例を見てみましょう。どのように処理が行われているかを理解するには、『NumPyの配列のスライスの必須テクニックまとめ』の理解が役立ちます。ぜひ確認しておいてください。

import numpy as np
a = np.array([0, 1]).reshape((2,1,1))
a

array([[[0]],

       [[1]]])

c1 = np.array([1, 2, 3]).reshape((1,3,1))
c1

array([[[1],
        [2],
        [3]]])

c2 = np.array([-1, -2, -3, -4, -5]).reshape((1,1,5))
c2

array([[[-1, -2, -3, -4, -5]]])

# 出力結果は2x3x5
# res[0,:,:] を c1に
# res[1,:,:] を c2 に
np.choose(a, (c1, c2))

array([[[ 1,  1,  1,  1,  1],
        [ 2,  2,  2,  2,  2],
        [ 3,  3,  3,  3,  3]],

       [[-1, -2, -3, -4, -5],
        [-1, -2, -3, -4, -5],
        [-1, -2, -3, -4, -5]]])

3. まとめ

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