NumPy

NumPy Quickstart チュートリアルに取り組む (5-1) 高度なインデックス作成とトリック インデックスの配列によるインデックス作成

NumPyは通常のPythonシーケンスよりも多くのインデックス機能を提供している。整数やスライスによるインデックス付けに加えて、配列は整数の配列やブーリアンの配列によるインデックス付けが可能。

import numpy as np
a = np.arange(12)**2
i = np.array([1, 1, 3, 8, 5])
print(a[i])  # [ 1  1  9 64 25]

j = np.array([[3, 4], [9, 7]])
print(a[j])
'''
[[  9 16 ]
 [ 81 49 ]]
'''

インデックス付き配列aが多次元の場合、インデックスの1つの配列はaの1次元目を参照する。以下の例では、ラベルの画像をパレットを用いてカラー画像に変換することでこの動作を示す。

import numpy as np
palette = np.array([[0, 0, 0],         # black
                    [255, 0, 0],       # red
                    [0, 255, 0],       # green
                    [0, 0, 255],       # blue
                    [255, 255, 255]])  # white
image = np.array([[0, 1, 2, 0],        # each value corresponds to a color in the palette
                  [0, 3, 4, 0]])
palette[image]                         # the (2, 4, 3) color image

'''
array([[[  0,   0,   0],
        [255,   0,   0],
        [  0, 255,   0],
        [  0,   0,   0]],

       [[  0,   0,   0],
        [  0,   0, 255],
        [255, 255, 255],
        [  0,   0,   0]]])
'''

複数の次元のインデックスを与えることもできる。各次元のインデックスの配列は同じ形状でなければならない。

import numpy as np
a = np.arange(12).reshape(3,4)
print(a)
'''
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
'''

i = np.array([[0, 1],                     # indices for the first dim of a
              [1, 2]])
j = np.array([[2, 1],                     # indices for the second dim
              [3, 3]])

print(a[i, j])                                   # i and j must have equal shape
'''
[[ 2  5]
 [ 7 11]]
'''

print(a[i, 2])
'''
[[ 2  6]
 [ 6 10]]
'''

print(a[:, j])                                     # i.e., a[ : , j]
'''
[[[ 2  1]
  [ 3  3]]

 [[ 6  5]
  [ 7  7]]

 [[10  9]
  [11 11]]]
'''

Pythonでは、arr[i, j]はarr[(i, j)]と全く同じのため、iとjをタプルに入れて、インデックスを作成することができる。

l = (i, j)
print(a[l])
'''
[[ 2  5]
 [ 7 11]]
'''

しかし、この配列はaの1次元目をインデックス化していると解釈されてしまうので、iとjを配列に入れることはできない。

s = np.array([i, j])

# print(a[s])
# IndexError: index 3 is out of bounds for axis 0 with size 3

print(a[tuple(s)])
'''
[[ 2  5]
 [ 7 11]]
'''