np.argsort()関数 (あるいはndarrayの.argsort()
メソッド)を用いると、配列の小さい順のインデックスを得ることができます。
python
1 arr = np . array ( [ 2 , 6 , 4 , 1 , 5 , 3 ] )
2 idx = np . argsort ( arr )
3
4 print ( idx ) # [3 0 5 2 4 1]
5 print ( arr [ idx ] ) # [1 2 3 4 5 6]
なお引数にaxis=None
を指定すると多次元配列を一次元配列とみなして処理します。
以下は.argsort()
メソッドで小さい順にn番目までのインデックスを取得し、一次元インデックスをnp.unravel_index()関数 を用いて二次元座標にする自作関数の例です(渡す配列が三次元以上でも機能します)。
python
1 def get_min_index ( arr , n ) :
2 return np . unravel_index ( np . argsort ( arr , axis = None ) [ : n ] , arr . shape )
なお、結果が小さい順に並んでいなくても良いなら、上記のnp.argsort(arr, axis=None)
をnp.argpartition(arr, n, axis=None)
に変更することで速くなります。(argpartition()
は、やっていることはargsort()
より単純なのですが言葉で説明すると複雑なので説明は割愛します……)
実行例:
python
1 In [ 11 ] : Mat = np . random . default_rng ( 0 ) . random ( ( 4 , 5 ) )
2 . . . : print ( Mat )
3 [ [ 0.63696169 0.26978671 0.04097352 0.01652764 0.81327024 ]
4 [ 0.91275558 0.60663578 0.72949656 0.54362499 0.93507242 ]
5 [ 0.81585355 0.0027385 0.85740428 0.03358558 0.72965545 ]
6 [ 0.17565562 0.86317892 0.54146122 0.29971189 0.42268722 ] ]
7
8 In [ 12 ] : idx = get_min_index ( Mat , 5 )
9 . . . : print ( idx )
10 . . . : print ( Mat [ idx ] )
11 ( array ( [ 2 , 0 , 2 , 0 , 3 ] , dtype = int64 ) , array ( [ 1 , 3 , 3 , 2 , 0 ] , dtype = int64 ) )
12 [ 0.0027385 0.01652764 0.03358558 0.04097352 0.17565562 ]
13
14
15 In [ 13 ] : Mat = np . array ( [ [ 36 , 41 , 60 , 43 , 51 , 36 , 40 , 53 ] ,
16 . . . : [ 3 , 77 , 29 , 65 , 32 , 49 , 73 , 71 ] ,
17 . . . : [ 39 , 18 , 7 , 2 , 53 , 24 , 47 , 55 ] ,
18 . . . : [ 78 , 69 , 69 , 47 , 67 , 69 , 50 , 2 ] ,
19 . . . : [ 16 , 60 , 73 , 38 , 20 , 70 , 30 , 67 ] ,
20 . . . : [ 8 , 40 , 20 , 35 , 78 , 73 , 44 , 15 ] ,
21 . . . : [ 40 , 6 , 16 , 29 , 17 , 59 , 26 , 65 ] ,
22 . . . : [ 42 , 9 , 48 , 71 , 42 , 27 , 45 , 16 ] ,
23 . . . : [ 7 , 69 , 76 , 7 , 34 , 67 , 79 , 46 ] ,
24 . . . : [ 51 , 29 , 19 , 70 , 52 , 54 , 34 , 45 ] ] )
25
26 In [ 14 ] : idx = get_min_index ( Mat , 10 )
27 . . . : print ( idx )
28 . . . : print ( Mat [ idx ] )
29 ( array ( [ 2 , 3 , 1 , 6 , 2 , 8 , 8 , 5 , 7 , 5 ] , dtype = int64 ) , array ( [ 3 , 7 , 0 , 1 , 2 , 3 , 0 , 0 , 1 , 7 ] , dtype = int64 ) )
30 [ 2 2 3 6 7 7 7 8 9 15 ]