Matlabのfind()をpythonで


遅いとか言われつつ、matlabではよく使うfind()
似たようなことはnumpy.argwhere()でできるっぽい。

できること

Matlabの

[n, m] = find(condition)

をpythonで書くと

import numpy as np

nm = np.argwhere(condition)
n = nm[:, 0]
m = nm[:, 1]

こんな感じになる。

例:Matlab find()

>> x=reshape(0:11, [], 4)';
>> [n, m] = find(mod(x, 3)==0)

n =

     1
     2
     3
     4


m =

     1
     1
     1
     1

例:python numpy.argwhere()

>>> import numpy as np
>>> x = np.arange(12).reshape(4, -1)
>>> np.argwhere(x % 3 == 0)

array([[0, 0],
       [1, 0],
       [2, 0],
       [3, 0]])

インデックスが0始まりか1始まりか、という差以外は同じ。

できない(?)こと

'first' 'last'

matlabの

find(condition, 1, 'first')
find(condition, 3, 'last')

最初から1つだけ、とか最後から3つだけ、みたいなのは、numpy.argwhere()だけではできないっぽい。
全部取ってきてから必要なところだけ拾うか、forで書き下すか?
なんか他に良いモジュール/メソッドがあるのかも。見つけたら更新します(たぶん)。

ベクトル展開した上でのインデックス

matlabの

n = find(condition)

のようなベクトルに展開した上でのインデックスをとってくるのはコマンド一発では無さそう(?)
個人的にはベクトル展開してインデックスをとってくるのはあまり直感的じゃないので、割とどうでもいい。
どうしても必要なら、行と列と行列サイズから簡単に求められるし。