在NumPy数组中查找给定的坐标位置

时间:2015-04-20 10:58:25

标签: python numpy scipy

import numpy as np

ref_cols = [11, 5, 12, 13, 15]
ref_rows = [1, 11, 2, 3, 5]
rows, cols = np.mgrid[1:6, 11:16]
print cols

[[11 12 13 14 15]
 [11 12 13 14 15]
 [11 12 13 14 15]
 [11 12 13 14 15]
 [11 12 13 14 15]]
print rows

[[1 1 1 1 1]
 [2 2 2 2 2]
 [3 3 3 3 3]
 [4 4 4 4 4]
 [5 5 5 5 5]]

我想得到给定的cols和行(11,1),(5,11),(12,2),(13,3),(15,5)存在的位置。所以预期答案如下:

[[True, False, False, False, False],
[False, True, False, False, False],
[False, False, True, False, False],
[False, False, False, False, False],
[False, False, False, False, True]]

我试过:

rows_indices = np.in1d(rows, ref_rows).reshape(rows.shape)
cols_indices = np.in1d(cols, ref_cols).reshape(cols.shape)
answers = (rows_indices & cols_indices)
print answers

但回答是错误的。

怎么做呢?

2 个答案:

答案 0 :(得分:2)

您尝试出错的原因是您首先需要单独评估每一对,并且您不能分别首先评估所有行和列,然后在逻辑操作中将它们组合。

以下是解决问题的一种方法:

out = np.zeros(rows.shape, dtype=bool)
for r, c in zip(ref_rows, ref_cols):
    out |= (r == rows) & (c == cols)
print out

答案 1 :(得分:2)

可能存在一个更加优雅的解决方案,但这对我有用,并以矢量化方式编写......

import numpy as np

ref_cols = np.array([11, 5, 12, 13, 15])
ref_rows = np.array([1, 11, 2, 3, 5])
rows, cols = np.mgrid[1:6, 11:16]

m = (cols[:,:,None]==ref_cols[None,None,:]) & (rows[:,:,None]==ref_rows[None,None,:])

answer = np.any(m,axis=2)
#array([[ True, False, False, False, False],
#       [False,  True, False, False, False],
#       [False, False,  True, False, False],
#       [False, False, False, False, False],
#       [False, False, False, False,  True]], dtype=bool)