通过另一个不同形状的数组过滤 numpy 数组答案

【问题标题】：filter numpy array by another array of different shape通过另一个不同形状的数组过滤 numpy 数组
【发布时间】：2016-04-13 09:31:24
【问题描述】：

给定：

a = [[0, 1], [2, 2], [4, 2]]
b = [[0, 1, 2, 3], [2, 2, 3, 4], [4, 2, 3, 3], [2, 3, 3, 3]]

解决办法是：

for (i,j) in zip(a[:, 0], a[:, 1]):
                print b[np.logical_and( a[:, 0] == i,  a[:, 1] == j)]

结果应该是

[[0 1 2 3]]
[[2 2 3 4]]
[[4 2 3 3]]

在不使用for-loop 的情况下是否有解决此问题的方法？

【问题讨论】：

标签： python arrays numpy vectorization

【解决方案1】：

您可以将NumPy broadcasting 用于矢量化解决方案，就像这样 -

# 2D mask corresponding to all iterations of : 
# "np.logical_and( a[:, 0] == i,  a[:, 1] == j)"
mask = (a[:,None,0] == a[:,0]) & (a[:,None,1] == a[:,1])

# Use column indices of valid ones for indexing into b for final output
_,C_idx = np.where(mask)
out = b[C_idx]

示例运行 -

In [67]: # Modified generic case
    ...: a = np.array([[0, 1], [3, 2], [3, 2]])
    ...: b = np.array([[0, 1, 2, 3], [2, 2, 3, 4], [4, 2, 3, 3], [2, 3, 3, 3]])
    ...: 
    ...: for (i,j) in zip(a[:, 0], a[:, 1]):
    ...:     print b[np.logical_and( a[:, 0] == i,  a[:, 1] == j)]
    ...:     
[[0 1 2 3]]
[[2 2 3 4]
 [4 2 3 3]]
[[2 2 3 4]
 [4 2 3 3]]

In [68]: mask = (a[:,None,0] == a[:,0]) & (a[:,None,1] == a[:,1])
    ...: _,C_idx = np.where(mask)
    ...: out = b[C_idx]
    ...: 

In [69]: out
Out[69]: 
array([[0, 1, 2, 3],
       [2, 2, 3, 4],
       [4, 2, 3, 3],
       [2, 2, 3, 4],
       [4, 2, 3, 3]])

【讨论】：

【解决方案2】：

给定：

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

计算：

temp = np.in1d(b[:,0], a[:,0]) * np.in1d(b[:,1], a[:,1])
result = b[temp]
print 'result:', result

输出：

result: [[0 1 2 3]
 [2 2 3 4]
 [4 2 3 3]]

【讨论】：