基於兩個numpy數組獲取排序索引的最有效方法

Question

我如何只考慮另一個numpy數組的某些索引（val）如何獲取numpy數組的排序索引（距離）。

例如，考慮下面的兩個numpy數組val和distance：

val = np.array([[10, 0,  0,  0,  0],
                [0,  0,  10, 0,  10],
                [0,  10, 10, 0,  0],
                [0,  0,  0,  10, 0],
                [0,  0,  0,  0,  0]])


distance = np.array([[4, 3, 2, 3, 4],
                     [3, 2, 1, 2, 3],
                     [2, 1, 0, 1, 2],
                     [3, 2, 1, 2, 3],
                     [4, 3, 2, 3, 4]])

val == 10的距離是4、1、3、1、0、2。我想將它們排序為0、1、1、2、3、4，然后從距離數組中返回相應的索引。

返回類似：

(array([2, 1, 2, 3, 1, 0], dtype=int64), array([2, 2, 1, 3, 4, 0], dtype=int64))

要么：

(array([2, 2, 1, 3, 1, 0], dtype=int64), array([2, 1, 2, 3, 4, 0], dtype=int64))

由於第二個和第三個元素都具有距離“ 1”，所以我猜索引可以互換。

嘗試使用np.where, np.argsort, np.argpartition, np.unravel_index組合np.where, np.argsort, np.argpartition, np.unravel_index但似乎無法正常工作

Answer 1

這是masking的一種方法-

In [20]: mask = val==10

In [21]: np.argwhere(mask)[distance[mask].argsort()]
Out[21]: 
array([[2, 2],
       [1, 2],
       [2, 1],
       [3, 3],
       [1, 4],
       [0, 0]])