numpy.where で配列から条件に一致するデータのインデックスを抽出する方法

numpy.where で配列条件に一致するインデックスを取得

NumPyでデータを操作する上で、ある条件に一致するデータのみを抽出するには numpy.where を使います。

numpy.where は、NumPy配列(ndarray)から条件に一致するデータのインデックスを抽出することが可能です。

numpy.where関数の使い方

numpy.where 関数は、引数で条件を指定すると、一致するデータのインデックスを返します。取得できるのは一致する要素そのものではなくインデックスです。

例えば、ある数列から偶数値のみを取り出すには次のようにします。

import numpy as np

data = np.arange(10, 20)
# [10 11 12 13 14 15 16 17 18 19]

# 偶数の要素のインデックスを取得
# (array([1, 3, 5, 7, 9], dtype=int64),)
index = np.where(data%2 == 1)

# インデックスを指定して要素を取得
# [11 13 15 17 19]
print(data[index])

numpy.where関数のパラメータ

numpy.where(condition[, x, y])

パラメータ 説明
condition array_like, bool 抽出する条件を指定します。
x, y array_like 省略可能。x, y両方が指定されていれば、trueの場合は x, falseの場合は yを返します。省略の場合は、一致する要素のインデックスを返します。

return値として条件に一致するインデックス、もしくは引数で指定したx, yを含む配列を返します。

サンプルコード

多次元配列から条件指定で抽出

上で1次元配列に対するコードを載せたので、ここでは多次元配列に対しての抽出を行います。

3列✕5行のデータを対象に、3列目のデータを条件に抽出を行います。

import numpy as np

# 3列✕5行の2次元配列
data = np.array([
    [1, 1, 0],
    [1, 0, 1],
    [0, 1, 0],
    [0, 0, 1],
    [1, 1, 1],
])

# 最終列の値が1のデータのみ取得
index = np.where(data[:, -1] == 1)

data[index]
# array([[1, 0, 1],
#        [0, 0, 1],
#        [1, 1, 1]])

条件の指定時に、列を指定して条件設定してやるとインデックスを抽出できます。

x, yを指定して取得結果を変換

パラメータの x, y をそれぞれ指定することで取得できる値をインデックスではなく、指定した x, y に変換できます。例えば奇数の値は "奇数" という文字列にし、偶数の値は "偶数" という文字列に変換した配列を取得できます。

import numpy as np

data = np.arange(1, 10)
# [1, 2, 3, 4, 5, 6, 7, 8, 9]

np.where(data%2 == 0, '偶数', '奇数')
# array(['奇数', '偶数', '奇数', '偶数', '奇数', '偶数', '奇数', '偶数', '奇数'], dtype='<U2')

1 の場合 True, 0 の場合 False とする場合を考えます。

import numpy as np

data = np.array([
    [1, 1, 0],
    [1, 0, 1],
    [0, 1, 0],
    [0, 0, 1],
    [1, 1, 1],
])

np.where(data == 1, True, False)
# array([[ True,  True, False],
#        [ True, False,  True],
#        [False,  True, False],
#        [False, False,  True],
#        [ True,  True,  True]])

以上。

参考URL