numpyでarray内要素の全組み合わせの計算をする

例えばたくさんの点の座標値が配列に入っていると考えてください。

この点同士の全組み合わせの距離を計算したい、あるいは一番近い点同士を見つけたいなどのタスクはnumpyだとちょっとだけめんどうです。毎回忘れるのでメモします。

>>> import numpy as np
>>> positions = np.array([[0, 0, 0], [5, 5, 5], [1, 2, 10], [-3, 0, -2]])
>>> positions.shape
(4, 3)

こんな感じで、xyz座標で表される3次元ユークリッド空間上の点が4つあるとしましょう。

この中で、一番距離が近い点の組み合わせはどれか？

それを知るためには、点同士の全組み合わせの距離を計算しなければなりません。素直にfor文で計算していたらマサカリが飛んでくるかサービスが完成しないでしょう。そして、そのためにnumpyは存在します。

以下のようにして、全組み合わせの距離を得ることができます。

>>> tmp_index = np.arange(positions.shape[0])
>>> xx, yy = np.meshgrid(tmp_index, tmp_index)
>>> distances = np.linalg.norm(positions[xx]-positions[yy], axis=2)
>>> distances
array([[ 0.        ,  8.66025404, 10.24695077,  3.60555128],
       [ 8.66025404,  0.        ,  7.07106781, 11.74734012],
       [10.24695077,  7.07106781,  0.        , 12.80624847],
       [ 3.60555128, 11.74734012, 12.80624847,  0.        ]])

新たに(4, 4)のarray distancesが得られました。distances[i][j]には、最初に定義したarray positionsのi番目の点とj番目の点の間の距離が格納されています。

meshgridから返ってきたxx, yyは同じ次元の配列になっていて、走査することでtmp_indexの全組み合わせを探索することができます。なんて便利な！！

>>> xx
array([[0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3],
       [0, 1, 2, 3]])
>>> yy
array([[0, 0, 0, 0],
       [1, 1, 1, 1],
       [2, 2, 2, 2],
       [3, 3, 3, 3]])

ここからは簡単ですね。同じ点同士を計算している対角成分をnanとすれば、nanminで最小値が取得できる状態になります。

>>> np.fill_diagonal(distances, np.nan)
>>> np.nanmin(distances)
3.605551275463989

ここから最小値を持つインデックスを求めるには、nanargminとunravel_indexを使うとできます。

Stack Overflow

Numpy minimum in (row, column) format

How can I know the (row, column) index of the minimum of a numpy array/matrix?For example, if A = array(, ]), I want to get (1, 1)Thanks!

>>> np.unravel_index(np.nanargmin(distances), distances.shape)
(0, 3)

0番目の点と3番目の点、すなわち点[0, 0, 0]と点[-3, 0, -2]の間の距離が3.605551275463989で最小であることがわかりました。

argminやnanargmin, argmax等を複数次元の配列に適用する方法は実は今回初めて知りました。unravel_index便利ですね。