技術メモ

役に立てる技術的な何か、時々自分用の覚書。幅広く興味があります。

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]])

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

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 ... 配列の中から一つ選び取ってくれる