PYTHON 行列から条件を満たす要素番号を探す

行列から条件を満たす要素番号を探す

numpyを使って、行列から条件を満たす要素番号を探す方法を記述します。

インポートするパッケージはnumpyです。

20行×3列のint型乱数行列を作成し、行列の各列から50を超える要素を where() を使って探してみます。

import numpy as np
# 乱数配列の作成
rng = np.random.default_rng()
# 20行×3列の乱数配列
arr = rng.random((20,3)) * 100
# 行列の型を変更
arr = arr.astype(np.int16)
print(arr)
# 行列1列目（X座標）に条件式>50 を適用し、条件を満たす要素番号を抽出
print(np.array(np.where(arr[::,0] > 50)))
# 行列2列目（Y座標）に条件式>50 を適用し、条件を満たす要素番号を抽出
print(np.array(np.where(arr[::,1] > 50)))
# 行列3列目（Z座標）に条件式>50 を適用し、条件を満たす要素番号を抽出
print(np.array(np.where(arr[::,2] > 50)))
# 行列1列目（X座標）に複数の条件式を適用したい場合
print(np.array(np.where((arr[::,0] >= 30) & (arr[::,0] > 50))))
# 条件に適合する要素がない場合
print(np.array(np.where(arr[::,0] > 100)))

[[ 3 60 87]
 [18 86 38]
 [98 36 42]
 [26 97 38]
 [49 71 60]
 [48 54 58]
 [72 53 40]
 [96 52 57]
 [40  3 27]
 [92 12  5]
 [13  8 24]
 [40 91 79]
 [63 36 63]
 [ 8 75 63]
 [43 50 58]
 [ 6  6  2]
 [89 94 30]
 [15 98 10]
 [54 78 55]
 [15  6 56]]
[[ 2  6  7  9 12 16 18]]
[[ 0  1  3  4  5  6  7 11 13 16 17 18]]
[[ 0  4  5  7 11 12 13 14 18 19]]
[[ 2  6  7  9 12 16 18]]
[]

実行すると、条件を満たす要素番号の配列が返ってきます。該当する条件が見つからない場合は空の配列が返ってきます。

次にwhere() を応用したデータ渡しの方法を説明します。

下のコードは、20行×3列のint型乱数行列を作成し、グラフに表示するコードです。

3列目の要素が50を超える要素を赤色で表示し、50以下の要素は緑色で表示しています。グラフへのデータ渡しにwhere() を使っています。

グラフをインタラクティブにする為、ipymplというパッケージのインストールが必要です。

$ pip install ipympl

%matplotlib widget
import numpy as np
import matplotlib.pyplot as plt
# 乱数配列の作成
rng = np.random.default_rng()
# 10行×3列の乱数配列
arr = rng.random((20,3)) * 100
# 行列の型を変更
arr = arr.astype(np.int16)
# 3列目の数値が50を超える要素を抽出
z_over50 = np.where(arr[::,2] > 50)
z_under50 = np.where(arr[::,2] <= 50)
# グラフを作成
fig = plt.figure(figsize=(8, 5))
ax = fig.add_subplot(projection='3d')
ax.set_xlabel('x', fontsize=8)
ax.set_ylabel('y', fontsize=8)
# グラフに乱数配列を渡す
ax.scatter(arr[::,0][z_over50], arr[::,1][z_over50], arr[::,2][z_over50], color='red', s=1, lw=1)
ax.scatter(arr[::,0][z_under50], arr[::,1][z_under50], arr[::,2][z_under50], color='green', s=1, lw=1)

表示されるグラフの角度をマウスで調整して確認すると、Z=50 を境に点の色が赤色に変わっている事が確認できます。

PYTHON 条件によるグラフの色分け