基于两个numpy数组获取排序索引的最有效方法

时间:2018-12-02 08:00:11

标签: python numpy

如何仅考虑来自另一个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,2,3,4并从distance返回相应的索引数组。

返回类似的内容

(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的组合,但似乎无法使其正常工作

1 个答案:

答案 0 :(得分:3)

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