Python Numpy通过搜索2个值来查找由5个值组成的数组的行索引

时间:2019-03-17 14:05:34

标签: python arrays numpy

我很抱歉,如果标题令人困惑,我是numpy的新手,不习惯使用该术语。

假设我们有一个numpy数组充当世界地图。 参数为(x y r g b)-全部为int16

示例:

a = np.array([[  0,   0,   0, 255,   0], #index 0
              [  0,   1,   0,   0, 255], #index 1
              [  0,   2,   0, 255,   0]]) #index 2

现在我们要查找具有x和y值(0,2)的行的索引值-因此找到具有索引2的行。

[  0,   1,   0,   0, 255] #index 2

如何在不输入其余值(r g b)的情况下执行此操作?基本上,我们正在搜索具有两个值的五值行-我该怎么做?

2 个答案:

答案 0 :(得分:1)

您可以将行切成第二列,并检查它们是否等于[0,2]。然后使用allaxis设置为1,将满足所有条件的参数设置为True,并使用布尔数组为ndarray编制索引:

a = np.array([[  0,   0,   0, 255,   0],
              [  0,   1,   0,   0, 255],
              [  0,   2,   0, 255,   0]])

a[(a[:,:2] == [0,2]).all(1)]
# array([[  0,   2,   0, 255,   0]])

答案 1 :(得分:0)

这是您的数据:

import numpy as np

arr = np.array([[  0,   0,   0, 255,   0],
              [  0,   1,   0,   0, 255],
              [  0,   2,   0, 255,   0]])

a,b = 0,2 # [a,b] is what we are looking for, in the first two cols

以下是获取包含[a,b]的行索引的解决方案:

found_index = np.argmax(np.logical_and(arr[:,0]==[a],arr[:,1]==[b]))
print (found_index)

输出:

  

2

说明:

了解其工作原理的最佳方法是打印其中的每个部分:

print (arr[:,0]==[a])

输出:

  

[True True True]

print (arr[:,1]==[b])

输出:

  

[False False True]

print (np.logical_and(arr[:,0]==[a],arr[:,1]==[b]))
# print (np.logical_and([ True  True  True], [False False  True]))

输出:

[False False True]