2次元配列の特定の条件を満たす要素の中からランダムに選択【numpy】

「0~1の値を取る2次元配列があって、0.7以上の要素を持つインデックスの中からランダムにインデックスを取ってくる。」
こんなことをしたい時にnumpyを駆使すれば綺麗に書けたのでメモしておく。

例えばこんな配列(numpyのarray)があったとする

import numpy as np
a = np.random.randint(0,10,size=[8,6])
array([[6, 9, 4, 1, 2, 4],
       [4, 9, 3, 8, 9, 7],
       [9, 6, 6, 8, 9, 7],
       [6, 6, 2, 7, 2, 3],
       [9, 9, 3, 2, 1, 0],
       [9, 9, 0, 3, 8, 2],
       [0, 5, 4, 4, 0, 7],
       [1, 8, 6, 0, 9, 4]])

この中から８以上の値を持つ要素は

a > 7
array([[False,  True, False, False, False, False],
       [False,  True, False,  True,  True, False],
       [ True, False, False,  True,  True, False],
       [False, False, False, False, False, False],
       [ True,  True, False, False, False, False],
       [ True,  True, False, False,  True, False],
       [False, False, False, False, False, False],
       [False,  True, False, False,  True, False]], dtype=bool)

これだけある。
このTrueの中からランダムに一つインデックスを選びたい。
例えば、[0,1]とかだ。

これをnumpyを使わずに書こうとするとこんな風になるだろう。

import random
index_list = []
for i in range(len(a)):
    for j in range(len(a[0])):
        if a[i][j] > 7:
            index_list.append([i,j])
coord = random.choice(index_list)

randomの便利な関数のおかげでこれでも割とすっきりしているが、
numpyを使えばさらにすっきり書ける。

index = np.random.choice(np.where(a.reshape(-1,) > 7)[0])
coord = [index/a.shape[1], index%a.shape[1]]

たったこれだけだ。
numpy凄い。

---

※解説

• reshape[-1,] ... 2次元配列を1次元配列に変換できる
• np.where ... Trueの箇所のインデックスを返してくれる
• np.random.choice ... 配列の中から一つ選び取ってくれる