二维 numpy 数组搜索（相当于 Matlab 的 intersect 'rows' 选项）

Question

我有两个4列的numpy数组（2D），每个数组都有几百行（都是浮点数），分别叫做cap和usp。现在我想关注每个数组中的3列子集（比如说capind=cap[:,:3]）：

1. 这两个数组之间有很多相同的行。
2. 每一行的组合在每个数组中都是唯一的。

我想找到一个高效的方法，来识别这两个数组中相同的三列（行）子集，同时保留每个数组的第4列，以便后续处理。简单来说，我想找到一个很好的numpy方法，来实现类似于Matlab的intersect函数的行选项（也就是([c, ia, ib]=intersect(capind, uspind, 'rows');)）。

这个函数会返回匹配行的索引，这样我就可以很简单地从原始数组中获取匹配的三元组和第4列的值（比如matchcap=cap[ia,:]）。

---

我目前的方法是基于论坛上一个类似问题的，因为我找不到完全匹配我问题的解决方案。不过，这个方法似乎有点低效，考虑到我的目标（我也还没有完全解决我的问题）：

这些数组大概是这样的：

cap=array([[  2.50000000e+01,   1.27000000e+02,   1.00000000e+00,
      9.81997200e-06],
   [  2.60000000e+01,   1.27000000e+02,   1.00000000e+00,
      9.14296800e+00],
   [  2.70000000e+01,   1.27000000e+02,   1.00000000e+00,
      2.30137100e-04],
   ...,
   [  6.10000000e+01,   1.80000000e+02,   1.06000000e+02,
      8.44939900e-03],
   [  6.20000000e+01,   1.80000000e+02,   1.06000000e+02,
      4.77729100e-03],
   [  6.30000000e+01,   1.80000000e+02,   1.06000000e+02,
      1.40343500e-03]])

usp=array([[  4.10000000e+01,   1.31000000e+02,   1.00000000e+00,
      5.24197200e-06],
   [  4.20000000e+01,   1.31000000e+02,   1.00000000e+00,
      8.39178800e-04],
   [  4.30000000e+01,   1.31000000e+02,   1.00000000e+00,
      1.20279900e+01],
   ...,
   [  4.70000000e+01,   1.80000000e+02,   1.06000000e+02,
      2.48667700e-02],
   [  4.80000000e+01,   1.80000000e+02,   1.06000000e+02,
      4.23304600e-03],
   [  4.90000000e+01,   1.80000000e+02,   1.06000000e+02,
      1.02051300e-03]])

然后我把每个4列的数组（usp和cap）转换成3列的数组（capind和uspind，下面用整数表示，方便查看）。

capind=array([[ 25, 127,   1],
   [ 26, 127,   1],
   [ 27, 127,   1],
   ...,
   [ 61, 180, 106],
   [ 62, 180, 106],
   [ 63, 180, 106]])
uspind=array([[ 41, 131,   1],
   [ 42, 131,   1],
   [ 43, 131,   1],
   ...,
   [ 47, 180, 106],
   [ 48, 180, 106],
   [ 49, 180, 106]])

使用集合操作可以让我找到匹配的三元组：carray=np.array([x for x in set(tuple(x) for x in capind) & set(tuple(x) for x in uspind)])。

这个方法似乎很好地找到了uspind和capind数组中共同的行值。接下来，我需要从匹配的行中获取第4列的值（也就是把carray和原始数组（cap和usp）的前三列进行比较，并以某种方式抓取第4列的值）。

有没有更好、更高效的方法来实现这个？如果有任何关于如何从源数组中提取第4列值的帮助，我将非常感激。

Answer 1

在Matlab中，和numpy一样用来返回行索引的方法如下，它会返回一个布尔数组，其中相同的行的索引会标记为1：

def find_rows_in_array(arr, rows):
    '''
    find indices of rows in array if they exist
    '''
    tmp = np.prod(np.swapaxes(
        arr[:, :, None], 1, 2) == rows, axis=2)
    return np.sum(np.cumsum(tmp, axis=0) * tmp == 1,
                  axis=1) > 0

上面的代码只返回非重复行的索引。如果你想返回所有可能的行索引，可以使用：

def find_rows_in_array(arr, rows):
    '''
    find indices of rows in array if they exist
    '''
    tmp = np.prod(np.swapaxes(
        arr[:, :, None], 1, 2) == rows, axis=2)
    return np.sum(tmp,
                  axis=1) > 0

这个方法要快很多。你可以互换数组作为输入，这样就能为每个数组找到对应的索引。祝你玩得开心 :D

Answer 2

numpy_indexed这个包（声明：我是它的作者）包含了你所需要的所有功能，而且实现得非常高效（也就是说，它是完全向量化的，所以在Python层面没有慢速循环）：

import numpy_indexed as npi
c = npi.intersection(capind, uspind)
ia = npi.indices(capind, c)
ib = npi.indices(uspind, c)

根据你对简洁性和性能的重视程度，你可能会更喜欢：

import numpy_indexed as npi
a = npi.as_index(capind)
b = npi.as_index(uspind)
c = npi.intersection(a, b)
ia = npi.indices(a, c)
ib = npi.indices(b, c)

Answer 3

根据你所假设的每个矩阵的行都是唯一的，并且有一些行是相同的，这里有一个解决方案。基本的思路是把两个数组合并在一起，然后排序，这样相似的行就会放在一起，接着对行之间进行比较。如果行是相同的，前面三个值应该接近零。

[原始]

## Concatenate the matrices together
cu = concatenate( (cap, usp), axis=0 )
print cu

## Sort it
cu.sort( axis=0 ) 
print cu

## Do a forward difference from row to row
cu_diff = diff( cu, n=1, axis=0 )

## Now calculate the sum of the first three columns 
##  as it should be zero (or near zero)
cu_diff_s = sum( abs( cu_diff[:,:-1] ), axis=1 ) 

## Find the indices where it is zero
##  Change this to be <= eps if you are using float numbers 
indices = find( cu_diff_s == 0 )
print indices

## And here are the rows...
print cu[indices,:]

我根据你上面的例子构造了一个数据集。看起来是有效的。可能有更快的方法，但这种方式不需要循环任何东西。（我不喜欢循环 :-)）。

[更新]

好的。我在每个矩阵中又添加了两列。倒数第二列在cap中是1，在usp中是2。最后一列只是原始矩阵的索引。

## Store more info in the array
##  The first 4 columns are the initial data
##  The fifth column is a code of 1 or 2 (ie cap or usp)
##  The sixth column is the index into the original matrix

cap_code = concatenate(  (ones( (cap.shape[0], 1 )), reshape( r_[0:cap.shape[0]], (cap.shape[0], 1))), axis=1 )
cap_info = concatenate( (cap, cap_code ), axis=1 )

usp_code = concatenate(  (2*ones( (usp.shape[0], 1 )), reshape( r_[0:usp.shape[0]], (usp.shape[0], 1))), axis=1 )
usp_info = concatenate( (usp, usp_code ), axis=1 )

## Concatenate the matrices together
cu = concatenate( (cap_info, usp_info), axis=0 )
print cu

## Sort it
cu.sort( axis=0 )
print cu

## Do a forward difference from row to row
cu_diff = diff( cu, n=1, axis=0 )

## Now calculate the sum of the first three columns 
##  as it should be zero (or near zero)
cu_diff_s = sum( abs( cu_diff[:,:3] ), axis=1 )

## Find the indices where it is zero
##  Change this to be <= eps if you are using float numbers 
indices = find( cu_diff_s == 0 )
print indices

## And here are the rows...
print cu[indices,:]
print cu[indices+1,:]

根据我构造的数据，这个方法似乎有效。虽然有点复杂，但我觉得我不想再继续这个方向了。

祝好运！

Answer 4

试试用字典。

capind = {tuple(row[:3]):row[3] for row in cap}
uspind = {tuple(row[:3]):row[3] for row in usp}

keys = capind.viewkeys() & uspind.viewkeys()
for key in keys:
    # capind[key] and uspind[key] are the fourth columns